Эффективное обучение на одной GPU
Рассмотрим способы улучшить эффективность обучения моделей на одной GPU.
Здесь https://teletype.in/@abletobetable/gpu_operations_and_memory более подробно про операции модели, которые выполняются на видеокарточке, а также про то, что влияет на утилизацию видеопамяти.
Можно почитать документацию huggingface: https://huggingface.co/docs/transformers/perf_train_gpu_one
При обучении нейросетей следим за 2мя ключевыми параметрами:
- Скорость обучения
Как быстро модель сходится, сколько данных успеваем показать, сколько денег тратим на вычисления - Качество модели
Валидационные и тестовые метрики, лоссы
То есть мы хотим как можно быстрее, дешевле и лучше (с точки зрения метрик) обучить модель.
Выбор batch size
Чтобы достичь наилучшей производительности, рекомендуется правильно подобрать размер батча.
Увеличиваем батч пока не OOM
Обычно нужно брать максимальный размер батча, который помещается на GPU, чтобы по полной утилизировать видеопамять.
Следует придерживаться стратегии выбора размера, который является степенью двойки: 2^N, потому что GPU быстрее обрабатывают тензоры, у которых “красивые” размерности из-за особенностей железа.
Как показывают nvidia в своем репорте про эффективные размеры батча вот тут - лучше выбирать batch size кратный 4 (при TF32) / 8 (при FP16) / 16 (при INT8), а для серии A100 - кратно 32 (TF32) / 64 (FP16) / 128 (INT8).
Искусственно увеличиваем размер или Gradient accumulation
Если необходимый размер батча не влезает в GPU, то можно использовать накопление градиентов для повышения фактического размера батча.
Например, на карточке есть место для одного батча, но делать градиентный шаг каждый семпл часто неоптимально, так как градиент будет шумным, вместо этого можно суммировать градиенты grad_accum шагов и после усреднения делать шаг оптимизации. Таким образом наш фактический размер батча равен gpu_batch_size * grad_accum.
Важно не забыть сделать именно усреднение, а не простое суммирование градиентов.
Накопление градиентов положительно влияет на сходимость, но на скорость обучения может сказываться негативно, так как дополнительные forward - backward будут замедлять обучение.
Например, если хочется использовать фактический batch size = 64, но влезает только 4, то лучше брать batch_size=4 + grad_accum=16, чем batch_size=1+grad_accum=64 из-за излишнего количества forward-backward и недостаточной утилизации видеопамяти.
Выбор числа шагов накопления градиентов стоит рассматривать только с точки зрения трейд оффа между эффективным размером батча и задержками на оптимизационные шаги.
Причем “красивое” число накопления градиентов не так важно, как в случае с размером батча, потому что с технической точки зрения, накопление градиентов не зависит от аппаратной структуры GPU.
Gradient Checkpointing
Иногда даже размер батча равный 1 может не влезть на видеокарту, потому что на GPU лежат не только градиенты и сами данные, но и различные активации. В таком случае можно применять gradient checkpointing.
Подробное объяснение того, как работает gradient checkpointing в этом блоге.
Но если коротко, то во время backward нам нужны активации слоев после forward:
- Default, GPU-rich метод
Можно хранить все активации на GPU, тогда делаем только один forward и быстро делаем backward, хотя тратим много GPU памяти.
- GPU-poor метод
Ни одну активацию не храним, поэтому экономим память GPU, но для backward-а делаем forward n-раз, чтобы пересчитать активации, что сильно замедляет обучение.
- Gradient checkpointing
Компромисс - сохраняем не все активации, а только самые необходимые, тем самым уменьшаем потребление видеопамяти с небольшой задержкой обучения.
Трейд офф между скоростью и потреблением GPU памятью:
Mixed precision training
Обычно модели хранят свои параметры в полной точности, то есть в fp32 - соответственно каждый параметр кодируем 32 битами: 1 бит на знак, 8 бит на экспоненту и 23 бита на мантиссу, но не всем параметрам нужна такая высокая точность, чтобы сохранить хорошее качество. Поэтому для ускорения вычислений и уменьшения потребления памяти (иногда), используют mixed precision.
Смысл mixed precision в том, что мы часть операций на GPU выполняем в полной точности (fp32), а часть, например, в половинной (fp16).
Операции в fp16 быстрее, хотя менее точные - но такой трейд офф обычно стоит того.
Варианты используемой точности во время обучения:
Mixed precision из коробки доступен во всех популярных фреймворках, например, в pytorch достаточно выполнять вычисления внутри контекстного менеджера autocast (дока и пример) - то есть самостоятельно переводить некоторые слои в режим fp16, а некоторые в fp32 не нужно.
Причем градиентный шаг делается в полной точности, поэтому нужен скейлер градиентов обратно в fp32 (пример в pytorch).
Стоит учитывать, что так как разные операции выполняются с разной точностью, то нужно хранить в памяти сразу 2 версии модели: в fp32 и в fp16, поэтому на один параметр модели будет нужно 4 + 2 = 6 байт. То есть мы ускоряем вычисления за счет более низкой точности в некоторых операциях и уменьшаем необходимый объем видеопамяти для обработки батча, хотя увеличиваем объем занимаемой GPU памяти для хранения весов модели, поэтому mixed precision эффективнее работает с небольшими моделями и большими батчами, а не наоборот, потому что при очень большой модели оверхед на хранение весов будет больше, чем экономия при вычислениях.
Оптимизации attention (fused kernels)
Известная проблема трансформеров заключается в том, что механизм self-attention растет квадратично по вычислительным ресурсам и памяти с увеличением числа входных токенов. Эта проблема еще более усугубляется в моделях, которые обрабатывают очень длинные последовательности.
Чтобы решить эту проблему, можно использовать техники оптимизации механизма внимания, например, FlashAttention2 или scaled dot product attention в PyTorch (SDPA), которые являются более эффективными с точки зрения памяти реализациями механизма внимания и могут ускорить вычисления.
Рекомендую к прочтению блогпост про оптимизацию работы модели на GPU с точки зрения скорости и объема вычислений, а также проблемы чтения и записи данных.
Устройство памяти
На иерархии пропускной способности и объема памяти видно, что самая большая CPU память - самая медленная. А самый небольшой кусок GPU - SRAM является самым эффективным - внутри него данные “летают”.
Fused kernels
Обычно вычисления происходят вот так:
- считываем данные из DRAM
- отправляем данные в HBM
- отправляем данные в SRAM и делаем вычисления
- результат отправляем в HBM
- результат + мб еще что-то отправляем в SRAM и делаем следующие вычисления
- повторяем шаги 3,4,5
У нас тут очень много чтений и записей между разными видами памяти - это сильно замедляет вычисления.
Чтобы решить эту проблему есть так называемые fused kernels, которые не делают избыточных чтений/записей.
Например, в pytorch (документация) есть torch.compile(model), который возвращает скомпилированную версию модели, что сильно оптимизирует вычисления.
FlashAttention (FA, 2nd version, 3rd version) - это тоже fused kernel.
Авторы переписали алгоритм вычисления attention, минимизировав read/writes и добавив лучшее распараллеливание.
Матрица внимания NxN нигде в алгоритме FA не материализуется, то есть она хранится частями и вычисления с ней происходят частями и параллельно. Также FA вообще получилось сфьюзить из-за softmax online trick вот отсюда - чтобы можно было считать softmax не сохраняя весь вектор логитов в одном месте, а считая его частями.
FA работает очень быстро при этом это точный алгоритм вычисления внимания. Поэтому везде, где железо позволяет использовать FA - по умолчанию используйте его.
Выбор оптимизатора
Наиболее часто используемым оптимизатором для обучения моделей на основе трансформеров является Adam или AdamW (Adam с weight decay). Adam достигает хорошей сходимости за счет хранения скользящего среднего предыдущих градиентов, однако это добавляет дополнительную нагрузку на память, пропорциональную (8 байт на параметр) числу параметров модели. Чтобы решить эту проблему, можно использовать альтернативные оптимизаторы.
- adamw_apex_fused
используется, например, через https://github.com/NVIDIA/apex на GPU - adafactor
не хранит скользящие средние для каждого элемента в матрицах весов, вместо этого аггрегирует - суммирует по столбцам или колонкам, что уменьшает потребление памяти (чуть больше 4 байтов на параметр).
Хотя сходится хуже, чем обычный AdamW. - Квантованный adam, например, 8-bit adam
Вместо аггрегации храним все скользящие средние, но в квантованном виде и приводим в полную точность только перед шагом оптимизации. Похоже на mixed precision обучение.
Предзагрузка данных
Одним из важных требований для достижения высокой скорости обучения является способность подавать данные GPU с максимальной скоростью, которую она может обработать. По умолчанию все операции выполняются в основном процессе, и он может не успевать читать данные с диска достаточно быстро, что создает bottleneck и приводит к недостаточной загрузке GPU.
Для ускорение доступа GPU к данным рекомендовано использовать:
- pinned memory (блог nvidia)
По умолчанию выделение памяти на хосте (CPU) является pageable. GPU не может напрямую обращаться к данным из pageable памяти хоста, поэтому, когда вызывается передача данных из pageable памяти хоста в память устройства, драйвер CUDA должен сначала выделить временный массив, закрепленный в памяти (pinned), скопировать данные хоста в этот закрепленный (pinned) массив, а затем передать данные из закрепленного массива в память устройства, как показано ниже.
Как видно на рисунке, закрепленная (pinned) память используется в качестве промежуточной области для передачи данных с устройства на хост. Мы можем избежать затрат на передачу данных между pageable и закрепленной (pinned) памятью, напрямую выделяя наши массивы на хосте в закрепленной (pinned) памяти.
- увеличение числа воркеров в дата лоадере
можно увеличить число воркеров для предобработки данных, чтобы предзагрузка данных на GPU была быстрее.
В целом можно запускать профайлер и смотреть, как долго GPU простаивает в ожидании данных и потом принимать соответствующие решения по оптимизации.
Другие оптимизации
Часть трюков, о которых точно стоит хотя бы упомянуть:
- DeepSpeed ZeRO
если модель не влезает на одну GPU даже с небольшим батчом
документация - Parameter-Efficient Fine Tuning (PEFT)
замораживаем все параметры модели и добавляем небольшую обучаемую добавку
документация - Mixture of Experts
позволяет ускорить обучение, хотя требует большого кол-ва GPU памяти