Gradient Checkpointing

Контрольная фиксация градиентов — это метод обучения, позволяющий экономить память: он сохраняет только выбранные промежуточные активации прямого прохода, а остальные пересчитывает во время обратного распространения ошибки. Этот метод также называют контрольной фиксацией активаций; он обменивает дополнительное вычисление на снижение пикового потребления памяти. Несмотря на название, метод фиксирует активации, а не градиенты параметров или файлы модели, поэтому особенно полезен, когда тензоры активаций не позволяют нейронной сети поместиться в доступную память GPU.

Как работает контрольная фиксация градиентов#

Во время стандартного прямого прохода нейронная сеть вычисляет промежуточные тензоры, называемые активациями. Система автоматического дифференцирования сохраняет активации, необходимые для последующего вычисления градиентов, как описано в механизмах autograd в PyTorch. Из-за большой глубины сети, крупных пакетов, изображений высокого разрешения и длинных входных последовательностей сохранённые тензоры могут занимать много памяти.

Контрольная фиксация градиентов делит сеть на сегменты:

1. Во время прямого прохода сохраняются входы или граничные активации выбранных сегментов.
2. Остальные промежуточные активации внутри этих сегментов удаляются.
3. Во время обратного прохода каждый сегмент с контрольной фиксацией повторно выполняет прямой проход, чтобы восстановить недостающие значения.
4. Восстановленные активации сразу используются для вычисления градиентов.

API PyTorch для контрольной фиксации активаций предоставляет эту возможность через torch.utils.checkpoint. Аналогичные концепции реализованы в контрольной фиксации градиентов и рематериализации в JAX и контрольной фиксации ленты в TensorFlow.

Расположение контрольных точек определяет компромисс. Обычно чем больше областей фиксируется, тем больше памяти экономится, но тем больше операций приходится выполнять повторно. Выборочные контрольные точки вокруг блоков с большим объёмом активаций могут обеспечить лучший баланс, чем повторный расчёт всей сети.

Компромиссы и связанные методы#

Контрольная фиксация градиентов обычно не меняет параметры модели, градиенты и состояния оптимизатора. Она в первую очередь экономит память, занятую активациями, а обучение замедляется из-за повторного выполнения некоторых вычислений прямого прохода. Точный результат зависит от архитектуры, границ контрольных точек, формы пакета и оборудования.

Этот метод отличается от нескольких связанных подходов:

• Накопление градиентов обрабатывает несколько микропакетов до обновления весов модели, создавая более крупный эффективный пакет без одновременной загрузки всех примеров. Контрольная фиксация градиентов вместо этого сокращает число активаций, сохраняемых для каждого микропакета.
• Смешанная точность использует типы данных с меньшей точностью для выбранных операций. Процесс автоматической смешанной точности в PyTorch может сократить расход памяти и ускорить совместимые операции, тогда как контрольная фиксация намеренно добавляет вычисления.
• Уменьшение размера пакета сокращает расход памяти, позволяя одновременно обрабатывать меньше примеров. Контрольная фиксация может сделать возможным использование большего пакета или изображений с более высоким разрешением.
• Контрольные точки обучения сохраняют веса и состояние оптимизатора для восстановления обучения или последующего инференса. В руководстве PyTorch по контрольным точкам модели описан этот механизм сохранения, не связанный с пересчётом активаций.

Код с контрольными точками должен давать согласованные результаты при исходном и повторном выполнении прямого прохода. Изменяемое состояние, переносы между устройствами или неконтролируемая случайность внутри такой области могут вызвать ошибки или неправильные градиенты. В трансформерах с архитектурой декодера во время обучения с контрольными точками также может потребоваться отключить кэш ключей и значений, поскольку кэшированное состояние инференса может конфликтовать с повторным построением графа прямого прохода.

Пример для PyTorch#

В следующем примере блок с высоким потреблением памяти фиксируется во время одного шага обучения:

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

torch.manual_seed(0)

block = nn.Sequential(
    nn.Linear(1024, 4096),
    nn.ReLU(),
    nn.Linear(4096, 1024),
)
optimizer = torch.optim.AdamW(block.parameters())
inputs = torch.randn(8, 1024, requires_grad=True)
targets = torch.zeros_like(inputs)

optimizer.zero_grad(set_to_none=True)
outputs = checkpoint(block, inputs, use_reentrant=False)
loss = nn.functional.mse_loss(outputs, targets)
loss.backward()
optimizer.step()

Сохраняются только входы блока и необходимые граничные данные; внутренние активации восстанавливаются во время loss.backward(). В реальных проектах сравнивай запуски с контрольной фиксацией и без неё, используя такие измерения, как пиковый объём выделенной памяти GPU в PyTorch.

Примеры применения в реальных условиях#

• Компьютерное зрение высокого разрешения: для медицинской сегментации, обнаружения объектов на аэрофотоснимках и промышленного контроля могут использоваться большие изображения, карты признаков которых занимают больше памяти, чем веса модели. Контрольная фиксация выбранных этапов базовой сети позволяет сохранить разрешение изображений или увеличить число примеров в пакете вместо агрессивного уменьшения входных изображений.

• Обучение трансформеров на длинных последовательностях: объём памяти для активаций быстро растёт с увеличением длины последовательности и числа слоёв. Повторное вычисление блоков трансформера позволяет обрабатывать на том же ускорителе более длинный контекст или крупные микропакеты. В руководстве NVIDIA по повторному вычислению активаций показаны полное и выборочное повторное вычисление для слоёв трансформера.

Практические рекомендации#

Используй контрольную фиксацию градиентов, если профилирование показывает, что большую часть памяти занимают активации, а не параметры или состояние оптимизатора. Начни с крупных повторяющихся блоков и проверь пиковое потребление памяти, время итерации и результаты валидации, прежде чем расширять охват.

В документированных конфигурациях обучения Ultralytics YOLO к основным средствам экономии памяти относятся автоматическая смешанная точность, размер изображения и физический размер пакета. В справочнике Ultralytics AutoBatch объясняется автоматический выбор размера пакета с учётом доступной памяти GPU. Если локального оборудования всё равно недостаточно, облачное обучение на Ultralytics Platform предоставляет настраиваемые облачные GPU для управляемых запусков обучения. Контрольная фиксация градиентов может дополнять эти средства, если для собственной архитектуры PyTorch требуется более точное управление памятью активаций.