Не знаю — нужно ли вступление к статье, посвящённой ускорению машинного обучения (Machine Learning, ML)?

Ускорение обучения моделей — это именно то, в чём нуждаются все ML‑инженеры. Более быстрое обучение модели означает ускорение экспериментов, что, в свою очередь, ведёт к ускорению выпуска новых версий программных продуктов. Кроме того — чем выше скорость обучения — тем меньше ресурсов нужно на каждую итерацию обучения модели. Поэтому предлагаю перейти сразу к делу.

Контейнеризация

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

Это обеспечивает полное единообразие окружения при отладке модели, при её профилировании и, в итоге, при её обучении. ML‑инженеру меньше всего нужно, например, оптимизировать часть кода, который, из‑за ускорения, обеспечиваемого python12, больше не является узким местом системы. То же самое можно сказать и о некоей ошибке, которая не воспроизводится на разных версиях CUDA.

В качестве отправной точки можно воспользоваться предварительно собранными образами от NVIDIA. В них уже установлены CUDA, PyTorch и другие популярные библиотеки.

Docker‑контейнер — это оптимальное решение проблем, которые обычно описывают так: «Вообще‑то на моей машине это работает. И у меня нет ни малейшего понятия о том, почему это не работает у тебя».

Освоение профилировщика PyTorch

Прежде чем что‑либо оптимизировать — нужно понять то, сколько времени уходит на выполнение тех или иных фрагментов кода. Профилировщик PyTorch — это почти универсальный инструмент для профилирования кода, ответственного за обучение моделей. Он может фиксировать следующие показатели:

• Данные о выполнении кода на CPU.

• Данные о выполнении кода в ядрах CUDA.

• Историю потребления памяти.

Это — всё, что нам нужно. И профилировщик PyTorch очень просто включить.

Для того чтобы приступить к записи событий, достаточно внедрить код, отвечающий за обучение, в контекст профилировщика:

import torch.autograd.profiler as profiler

with profiler.profile(
  activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
  on_trace_ready=torch.profiler.tensorboard_trace_handler('./logs'),
) as prof:
  train(args)

После этого можно запустить TensorBoard и посмотреть записи результатов профилирования. Тут ещё надо не забыть установить TensorBoard‑плагин torch‑tb‑profiler.

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

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

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

with profiler.record_function("forward_pass"):
  result = model(**batch)

with profiler.record_function("train_step"):
  step(**result)

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

with profiler.record_function("transformer_layer:self_attention"):
  data = self.self_attention(**data)

...

with profiler.record_function("transformer_layer:encoder_attention"):
  data = self.encoder_attention(**data, **encoder_data)

Интерпретация записей трассировщика PyTorch

После того, как записи трассировщика PyTorch собраны, откроем их в TensorBoard. Вот как выглядят результаты профилирования CPU и CUDA:

Сразу же найдём основные показатели, характерные для любого сеанса обучения:

• Загрузка данных.

• Прямой проход.

• Обратный проход.

Обратный проход обрабатывается PyTorch в отдельном потоке (в нашем случае это — поток 16893), поэтому то, что к нему относится, легко идентифицировать.

Загрузка данных

Нам нужно, чтобы время загрузки данных стремилось бы к нулю.

И никаких компромиссов.

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

Для того чтобы увидеть области простоя GPU — достаточно взглянуть на показатели GPU Est. SM Efficiency и GPU Utilization в записях профилировщика. Области, в которых нет активности — это наши, так сказать, пациенты. Именно там GPU ничем не занят.

Вот пара простых подходов к решению этой проблемы:

• Обработка данных в фоновых процессах (без GIL).

• Выполнение аугментации и трансформации данных в параллельных процессах.

Если применяется DataLoader PyTorch — этого легко достичь, задав параметр num_workers. Ситуация осложняется в том случае, если используется IterableDataset, так как при таком подходе имеет место дупликация данных. Но и эту проблему можно решить, прибегнув к get_worker_info(). Речь идёт о такой настройке итераций, когда каждый воркер получает различные, непересекающиеся строки.

Для организации работы с возможностью более тонкой её настройки можно подумать о самостоятельной реализации многопроцессных платформ с применением multiprocessing.

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

Заводим дружбу с распределителем памяти

Кеширующий распределитель памяти PyTorch для CUDA — это сущность, с которой, определённо, стоит подружиться.

Когда выделяют память для тензоров, используя PyTorch и работая с CUDA‑устройством, PyTorch применяет кеширующий аллокатор памяти. Это так из‑за того, что cudaMalloc/ cudaFree — это «дорогие» операции, которых мы стремимся избегать. Поэтому в PyTorch имеется механизм распределения памяти, который стремится повторно использовать фрагменты памяти, ранее выделенные с помощью cudaMalloc. В результате, если в распоряжении аллокатора PyTorch имеются подходящие блоки памяти — он сразу же, не вызывая cudaMalloc, передаст их в наше распоряжение. При таком подходе cudaMalloc вызывается лишь в начале работы программы.

Но если работают с данными переменной длины — разные прямые проходы потребуют применения промежуточных тензоров разных размеров. В результате у аллокатора PyTorch может не оказаться подходящего блока памяти. В таком случае аллокатор впадает в панику и освобождает ранее выделенные блоки памяти, вызывая cudaFree. Освобождённую память он сможет использовать для выполнения новых операций по выделению памяти.

После этого аллокатор снова начинает создавать кеш, выполняя множество «тяжёлых» операций cudaMalloc. Обнаружить эту проблему можно, взглянув в раздел профилирования памяти TensorBoard, где отображаются записи профилировщика.

Эту проблему можно увидеть и в самих записях профилировщика. Она выглядит как вызовы cudaMalloc и cudaFree.

Здесь красной линией показана память, зарезервированная аллокатором, объём которой, как видно, постоянно меняется. Это означает, что аллокатор не в состоянии эффективно обрабатывать запросы на выделение памяти.

Если же выделение памяти не заставляет аллокатор паниковать — красная линия будет абсолютно прямой.

Как уже было сказано — подобное случается при использовании тензоров разных форм. Как это исправить?

Расширяемые сегменты памяти

Первое, что в такой ситуации стоит попробовать — это переключение аллокатора PyTorch в сравнительно новый режим:

PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"

Если параметр expandable_segments установлен в True — это указывает аллокатору на то, что он должен выделять память для CUDA так, чтобы эта память потом могла бы быть расширена. Это позволит лучше организовать работу в ситуациях, когда программа часто требует изменения размера выделенной памяти. Например — при изменении размеров пакетов данных.

В результате можно сказать о том, что этот параметр указывает аллокатору PyTorch на то, что выделенные блоки памяти, при необходимости, могут быть расширены. А это — именно то, что нам нужно. Хотя, если вариации размеров данных слишком велики, этот подход может и не привести к решению проблемы. Если так и случилось — имеет смысл рассмотреть ещё один подход.

Сокращение разброса изменения объёмов выделяемой памяти

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

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

Подробности об аллокаторе памяти PyTorch можно найти в этой статье.

Приведение в порядок истории выделения памяти

Мы стремимся к тому, чтобы использовать всю доступную память GPU. Это позволяет загружать большие пакеты данных и быстрее эти данные обрабатывать. Но при увеличении пакета данных рано или поздно можно столкнуться с ошибкой нехватки памяти. Что вызывает эту ошибку?

Тут, в отладочных целях, имеет смысл просмотреть историю выделения памяти. Записать её можно с помощью PyTorch, а визуализировать с помощью memory_viz.

• Запуск: torch.cuda.memory._record_memory_history(max_entries=100 000)

• Сохранение: torch.cuda.memory._dump_snapshot(file_name)

• Остановка: torch.cuda.memory._record_memory_history(enabled=None)

Визуализация данных о выделении памяти может выглядеть примерно так:

Ось X — это время, ось Y — это общая использованная память, а цветные блоки — это тензоры. В результате тут видны моменты выделения и освобождения памяти.

Тут можно заметить узкие пики — это короткоживущие тензоры, которые занимают много памяти. Щёлкнув по тензору, можно получить сведения о том, где именно была выделена память под этот тензор. Нам нужно минимизировать эти пики, так как они ограничивают эффективное использование памяти. Можно узнать о том, что именно их вызывает, и рассмотреть возможность использования другого способа вычисления того, что нужно вычислить.

Помимо пиков тут легко выявлять утечки памяти.

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

Вот полезная статья о работе с памятью GPU.

Ускорение модели и использование меньшего объёма памяти

Что может быть лучше того, что описано в предыдущем заголовке? Достичь этого можно, воспользовавшись ядром FlashAttention для вычисления скалярного произведения в блоках внимания.

Если вы об этом не слышали — то знайте, что это — подход к вычислению точного скалярного произведения в блоке внимания без явного конструирования матрицы внимания. Это оптимизирует операции, выполняемые GPU, что улучшает скорость работы системы, и, кроме того, серьёзно минимизирует объём потребляемой памяти. Пожалуй, у нас просто нет причин этим не воспользоваться.

К сожалению, одна причина не пользоваться этим замечательным инструментом, всё же, есть. Это — аппаратное обеспечение. FlashAttention работает только с числами точности fp16 и bf16 на совместимом «железе». А это — NVIDIA Ampere, Hooper и так далее.

Существуют и другие библиотеки, использующие механизм Flash Attention. Поэтому, при необходимости, можно попытаться подобрать что‑то другое, лучше подходящее к конкретной кодовой базе. Например — XFormers или Transformer Engine.

А ещё можно взять и воспользоваться возможностями, которые есть в самом фреймворке PyTorch! Дело в том, что новые версии PyTorch могут использовать механизм Flash Attention там, где он применим. Для включения этого режима нужно выполнять блоки внимания в менеджере контекста, который указывает то, какую именно стратегию внимания использовать.

Оптимизация обучения в системах с несколькими GPU и с избыточностью данных — FSDP

Если для обучения моделей используют несколько GPU, то самым простым подходом в этой ситуации будет применение класса DistributedDataParallel (DDP). В результате будут созданы несколько идентичных процессов, а градиенты будут агрегироваться на шаге обратного прохода.

Но такой подход неоптимален!

Проблема заключается в следующем: когда запускают одинаковые процессы — на каждом GPU работа ведётся с одинаковыми моделями и с одинаковыми состояниями оптимизатора. Это избыточно. Решение заключается в шардинге (распределении) данных на разных GPU. Сделать это можно с помощью PyTorhc-обёртки Fully Sharded Data Parallel (FSDP, параллелизм с полным шардингом данных).

Как работает этот механизм?

Мы уже говорили о том, что, при обучении модели на нескольких GPU с применением DDP, в распоряжении каждого процесса имеются одни и те же данные. Улучшить ситуацию можно, прибегнув к нескольким оптимизациям.

Шардинг состояния оптимизатора (ZeRO1)

При обучении модели с помощью DDP каждый процесс хранит полную копию состояний оптимизатора. При применении ZeRO1 состояния оптимизатора распределяются по всем рангам — так, что каждый из них хранит лишь некую порцию состояний оптимизатора. При выполнении обратного прохода каждый ранг нуждается в сборе лишь тех состояний оптимизатора, которые имеют отношение к его параметрам, необходимым для проведения шага оптимизации. Такое сокращение избыточности помогает беречь память.

В случае с Adam, когда размер хранимых параметров примерно в два раза превышает размер модели, шардинг состояния оптимизатора по 8 рангам означает, что один ранг хранит лишь четверть (2/8) общего размера состояния.

Шардинг градиентов (ZeRO2)

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

• Агрегируем все градиенты, имеющие отношение к состояниям, которые имеются у ранга.

• Проводим вычисления для выполнения шага оптимизации.

• Отправляем данные шага оптимизации для порции параметров всем остальным рангам.

Как можно заметить — теперь каждому рангу не нужно хранить полную копию градиентов. Градиенты можно отправлять соответствующему рангу в тот момент, как они окажутся доступными. В результате это даёт возможность ещё сильнее снизить пиковое потребление памяти.

Шардинг параметров модели (ZeRO3)

А это — без преувеличения — настоящий прорыв.

Зачем хранить полную копию модели в каждом ранге? Прибегнем к распределению параметров по всем рангам. Это позволит точно вовремя загружать нужные параметры при выполнении прямого и обратного проходов.

Если речь идёт о больших моделях — эти оптимизации могут очень серьёзно снизить потребление памяти.

Как пользоваться FSDP?

На самом деле, пользоваться FSDP довольно просто. Достаточно обернуть модель в FSDP:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP


model = FSDP(model)

# крайне важно получить параметры из обёрнутой модели
# так как возвращается лишь их порция (распределённая часть)
optimizer = optim.Adam(model.parameters())

# проводим обучение так, как в обычных условиях
train(model, optimizer)

Ещё можно задать стратегию шардинга, применяемую FSDP. Например, можно выбрать стратегию SHARD_GRAD_OP, что позволит достичь поведения системы, похожего на её поведение при применении оптимизации ZeRO2. Подробности о других стратегиях можно найти здесь.

Кроме того, оборачивать в FSDP можно подмодули. В вышеприведённом примере используется лишь модуль FSDP, что снизит эффективность вычислений и эффективность использования памяти. Работает это так: представим, что модель содержит 100 линейных слоёв. Если прибегнуть к конструкции FSDP(model), то будет создан лишь один модуль FSDP, в который будет обёрнута вся модель. В данном случае AllGather будет собирать полный набор параметров для всех 100 линейных слоёв, а значит — не будет сэкономлена память CUDA для шардинга параметров.

Можно явным образом создавать обёртки для подмодулей или задать политику их автоматического оборачивания. Подробности о FSDP можно узнать из этого руководства.

Волшебное ускорение кода с помощью torch.compile

Для того чтобы ускорить работу модели на несколько процентов — достаточно просто включить torch.compile. Вот и вся магия.

Torch наблюдает за графом выполнения кода и пытается скомпилировать его, преобразовав так, чтобы модель могла бы выполняться эффективно, почти без обращений к Python.

Самый простой способ использования компилятора — обернуть модель в torch.compile:

import torch

model = torch.compile(model)

Такая конструкция выполняется почти мгновенно. Реальная трассировка производится только при первом прямом проходе.

У torch.compile есть множество опций, которые стоит попробовать.

Компилятор torch — это большая тема, которую я собираюсь раскрыть в будущих статьях.

Здесь можно найти руководство по torch.compile.

Итоги

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