Что такое gradient checkpointing: память при обучении сокращают за счёт повторного вычисления активаций. Приём работает так: обычный forward кэширует промежуточные тензоры для backward, а checkpoint-режим сохраняет только входы сегментов и пересчитывает внутренние активации во время обратного прохода. Документация PyTorch описывает это как обмен вычислений на память.
Активации, параметры, градиенты и optimizer state
Память обучения складывается из нескольких независимых частей, и путать их нельзя. Документация Hugging Face по анатомии памяти GPU перечисляет веса модели, состояния оптимизатора, градиенты, активации и временные тензоры операций. Веса — это параметры, которые вы учите. Градиенты — производные функции потерь по этим параметрам, они нужны оптимизатору на каждом шаге. Optimizer state — дополнительные буферы, которые ведёт сам оптимизатор: например, моменты в Adam. Активации — промежуточные результаты forward-прохода, которые backward использует для вычисления производных.
Gradient checkpointing работает именно с активациями. Он не уменьшает число параметров, не сжимает optimizer state и не меняет размер градиентов. Документация Hugging Face по анатомии памяти отмечает, что размер активаций зависит от размера батча, длины последовательности, глубины модели и скрытого размера. Временные тензоры, которые создают операции вроде softmax и матричных умножений, освобождаются после каждой операции, и на них приём тоже не влияет напрямую.
Разделение важно для диагностики. Если обучение упирается в память, сначала стоит понять, какая часть её занимает. При длинных последовательностях и глубоких моделях активации становятся заметной долей общего расхода — об этом пишет блог PyTorch, называя память активаций растущим вкладом в общее потребление. Тогда checkpointing даёт эффект. Если основную память съедают веса и состояния оптимизатора, выигрыш будет скромнее, и решать задачу придётся другими средствами.
Что сохраняет обычный forward для backward
Обычный forward-проход считает выход слоя за слоем и складывает промежуточные результаты в память. Документация Hugging Face по gradient checkpointing формулирует это так: forward обычно кэширует все промежуточные активации, чтобы backward мог их переиспользовать. Документация по анатомии памяти добавляет, что активации вычисляются в прямом проходе и кэшируются для обратного, где по ним считаются градиенты.
Зачем это нужно. Backward применяет правило цепочки: производная сложной функции равна произведению производных её частей. Чтобы посчитать производную по входу слоя, нужно знать локальную производную самого слоя, а она часто зависит от значения выхода. Если выход не сохранён, его придётся получать заново. Обычный режим выбирает хранение: память тратится, зато обратный проход идёт без лишних вычислений.
Объём активаций зависит от архитектуры, размера батча и длины последовательности. Например, Hugging Face указывает квадратный рост матриц attention по длине последовательности. PyTorch описывает освобождение сохранённой активации после её использования в backward. Checkpointing сокращает набор сохраняемых промежуточных значений и восстанавливает нужные значения вычислением.
Полная трасса x=2, a=3, b=4 с производными
Разберём учебную цепочку из двух умножений. Пусть u = a·x, y = b·u, а функция потерь L = y²/2. Входы: x = 2, a = 3, b = 4. Это условный пример для проверки арифметики. Измерения на GPU для него не проводились.
Прямой проход:
- u = a·x = 3·2 = 6
- y = b·u = 4·6 = 24
- L = y²/2 = 24²/2 = 576/2 = 288
Обратный проход по правилу цепочки:
- dL/dy = y = 24
- dL/db = dL/dy · u = 24·6 = 144
- dL/du = dL/dy · b = 24·4 = 96
- dL/da = dL/du · x = 96·2 = 192
- dL/dx = dL/du · a = 96·3 = 288
Проверим согласованность: dL/dx = 288, dL/da = 192, dL/db = 144. Эти значения нужны как эталон. В этой детерминированной учебной цепочке оба рассматриваемых режима должны дать одинаковые производные. Ниже сравним, что сохраняется в памяти в обычном режиме и в checkpoint-варианте при одинаковом результате.
Повторное вычисление u в checkpoint-сегменте
Для примера объединим оба умножения в checkpoint-сегмент F(x, a, b) = b·(a·x) с выходом y. Тогда u = a·x находится внутри сегмента, и его можно восстановить при backward. Вход x и параметры a, b остаются доступны, а выход y нужен внешней функции потерь. Если сегмент ограничить только u = a·x, внешнее умножение b·u должно сохранить u для градиента по b; такое разбиение не иллюстрирует заявленное освобождение u.
Во время backward восстанавливается внутреннее u = a·x = 3·2 = 6. Оно нужно для dL/db = 24·6 = 144. Остальная цепочка даёт dL/du = 96, dL/da = 192 и dL/dx = 288. Восстанавливаются нужные внутренние значения сегмента; результат совпадает с полной трассой.
Сравним два режима в таблице.
| Шаг | Обычный forward/backward | Checkpoint-сегмент |
|---|---|---|
| Forward | вычислить y = b·(a·x), сохранить внутреннее u = 6 | вычислить тот же y, внутреннее u не сохранять |
| Граница сегмента | вход x, параметры a и b, выход y | те же вход, параметры и выход |
| Backward | взять сохранённое u для dL/db | восстановить u = a·x = 6 внутри сегмента |
| dL/db | 144 | 144 |
| dL/da | 192 | 192 |
| dL/dx | 288 | 288 |
| Дополнительная работа | повторный forward не нужен | внутренние активации восстанавливаются вычислением |
Таблица сравнивает удержание внутреннего u; фактический расход памяти библиотек здесь не измерялся.
Документация PyTorch описывает приём как обмен вычислений на память. Документация TensorFlow для tf.recompute_grad формулирует ту же идею: повторное вычисление функции при обратном распространении выполняет избыточный расчёт и снижает общий расход памяти ленты авто-дифференцирования. Оба описания говорят об одном механизме, но относятся к своим библиотекам.
Обмен вычислений на память вместо уменьшения модели
Checkpointing не делает модель меньше. Число параметров остаётся прежним, optimizer state тоже, размер градиентов не меняется. Меняется только то, сколько промежуточных активаций удерживается в памяти между проходами. Документация Hugging Face по gradient checkpointing отмечает: обучение обычно идёт медленнее, потому что backward пересчитывает отброшенные активации, зато checkpointing снижает память под активации.
Отсюда практическое следствие. Приём помогает, когда упор в память дают именно активации, и стоит вычислений. Блог PyTorch пишет, что относительно eager-режима приём уменьшает объём памяти, сохраняемой для backward, и добавляет стоимость вычислений из-за пересчёта. Конкретные проценты экономии и замедления зависят от архитектуры, размера батча, длины последовательности и стратегии разбиения на сегменты, поэтому универсальных чисел здесь нет.
Порядок действий при нехватке памяти:
- Оцените, какая часть памяти уходит на активации, а какая — на веса и optimizer state.
- Если активации заметны, включите checkpointing и замерьте память и время шага.
- Сравните выигрыш по памяти с потерями по скорости.
- Если памяти всё равно мало, рассмотрите разбиение на сегменты мельче или другие приёмы.
TrainingArguments.gradient_checkpointing и PyTorch use_reentrant
Два уровня настройки решают разные задачи. Документация Hugging Face по gradient checkpointing указывает: чтобы включить приём, установите gradient_checkpointing=True. Это флаг уровня Trainer, он включает механизм для модели целиком.
PyTorch даёт более низкий уровень — torch.utils.checkpoint. Документация PyTorch рекомендует использовать use_reentrant=False и указывает, что параметр use_reentrant следует передавать явно. Явное указание снимает неоднозначность выбора реализации.
| Уровень | Что задаёт | Где описано |
|---|---|---|
| TrainingArguments.gradient_checkpointing | включает checkpointing для обучения через Trainer | документация Hugging Face |
| torch.utils.checkpoint | выбор реализации и гранулярность сегментов | документация PyTorch |
| use_reentrant | явный выбор реализации checkpoint | документация PyTorch |
Не переносите детали одного уровня на другой. Флаг Trainer не заменяет настройку checkpoint-функции, а параметры PyTorch не описывают поведение Trainer. Документация Hugging Face относится к текущей ветке main, поэтому конкретный набор доступных аргументов стоит сверять с версией вашей библиотеки.
Риски случайности и разного поведения функции
Checkpointing повторно выполняет вычисления во время backward. PyTorch предупреждает о неверных градиентах, если функция между проходами меняется, например из-за глобального состояния. Для случайных операций PyTorch по умолчанию сохраняет и восстанавливает состояние RNG при preserve_rng_state=True. Отключение этого механизма или перенос тензоров на новое устройство внутри функции ограничивает эквивалентность с обычным проходом.
Практические меры:
- проверяйте настройку preserve_rng_state и охват используемых устройств.
- проверяйте, что функция детерминирована при одинаковом входе;
- сравнивайте градиенты checkpoint-режима и обычного на учебной трассе;
- следите за состоянием, которое функция меняет снаружи.
Численные эффекты заранее не обещаются: расхождение либо есть, либо нет, и это выясняется сравнением. В нашем примере производные совпали, потому что функция детерминирована. Для реальных слоёв с dropout или случайной маской проверка обязательна.
Частые вопросы
Что такое gradient checkpointing простыми словами?
Это режим обучения, при котором forward сохраняет только входы сегментов сети, а внутренние активации backward пересчитывает заново. Память под активации снижается, вычислений становится больше. Документация PyTorch называет это обменом вычислений на память.
Как включить gradient checkpointing в Hugging Face?
Документация Hugging Face указывает: установите gradient_checkpointing=True в TrainingArguments. Флаг включает приём для обучения через Trainer. Конкретные дополнительные аргументы сверяйте с версией библиотеки.
Что делает use_reentrant в PyTorch?
Это параметр checkpoint-функции, выбирающий реализацию механизма. Документация PyTorch рекомендует use_reentrant=False и советует передавать параметр явно. Явное значение снимает неоднозначность.
Почему обучение становится медленнее?
Backward повторно выполняет forward для сегментов, чтобы восстановить отброшенные активации. Документация Hugging Face отмечает, что обучение обычно идёт медленнее, а память под активации снижается. Величина замедления зависит от архитектуры и разбиения.
Чем checkpointing отличается от сохранения модели на диск?
При сохранении на диск записывают параметры, например a = 3 и b = 4, чтобы возобновить обучение позже. При activation checkpointing речь о промежуточном значении u = 6 внутри одного шага. Это разные объекты с одинаковым словом checkpoint.