DeepSpeed настройка: почему обучение упирается в память
Обучение больших моделей упирается не в скорость ускорителей, а в объём памяти. Веса, градиенты и состояния оптимизатора находятся на устройстве одновременно, причём оптимизатор занимает кратно больше, чем сама модель. Нехватка памяти проявляется не на старте, а на середине прогона — когда уже потрачены часы на подготовку данных и разогрев.
Цена такого сбоя выше, чем кажется. Оборудование простаивает между перезапусками, прогресс откатывается до последней контрольной точки, а часть экспериментов приходится урезать: уменьшать размер пакета, сокращать длину контекста, брать модель слабее нужной. Падение процесса обучения тормозит не одну итерацию, а всё направление — от исследования до вывода модели в продукт.
Настройка DeepSpeed — это управление тем, где именно живут состояния. Стадии ZeRO определяют, что шардируется между ускорителями, выгрузка переносит часть состояния в память хоста, а профиль конфигурации связывает выбранную схему с конкретным железом и топологией обмена. Ниже разбираем, как это закрывается на практике — от выбора стадии до стабильного прогона без падений и ручных перезапусков.
Что меняется после грамотной конфигурации обучения
Без настройки распределённого обучения узкое место почти всегда одно: состояния оптимизатора и активации не помещаются в память ускорителя, поэтому инженер режет размер порции и теряет загрузку. Устройства простаивают в ожидании синхронизации, и прогон растягивается. Грамотная конфигурация разводит состояния по устройствам и снимает этот затор — та же модель обучается в разы быстрее.
DeepSpeed даёт для этого набор режимов: разделение состояний оптимизатора и градиентов по устройствам, выгрузку весов и активаций в память хоста, накопление градиентов мелкими порциями. Смысл не во включении всего сразу, а в подборе комбинации под конкретную модель и конфигурацию узлов: избыточное разделение добавляет обмен между устройствами и съедает выигрыш, недостаточное — возвращает нехватку памяти.
После подбора режимов загрузка оборудования выходит на ровные 85–95% вместо провалов на синхронизации, а скорость обучения перестаёт зависеть от длины последовательности. Устойчивость обеспечивают контрольные точки: сохраняем состояние оптимизатора и планировщика, а не только веса, и пишем их асинхронно, чтобы запись не тормозила счёт.
Время восстановления после отказа узла падает с часов до минут — процесс поднимается с последней контрольной точки, а не с нуля. В сумме это даёт высокую долю прогонов, доходящих до целевой метрики без ручного вмешательства, и предсказуемые сроки: план обучения строится по измеренной скорости шага, а не по оптимистичной оценке.
DeepSpeed настройка: режимы ZeRO и выгрузка памяти
Выбор режима определяется тем, какие состояния обучения остаются на каждом ускорителе. Состояния оптимизатора занимают в разы больше места, чем сами параметры, следом идут градиенты, и только потом веса. Режимы ZeRO отличаются тем, какую из этих групп шардируют между устройствами, а выгрузка памяти добавляет второй уровень — в память процессора или на локальный накопитель.
| Вариант | Когда подходит | Ограничения |
|---|---|---|
| ZeRO-1: шардируются состояния оптимизатора | Модель уже помещается в память ускорителей, но хочется поднять размер пакета данных | Параметры и градиенты дублируются на каждом устройстве — выигрыш по памяти умеренный |
| ZeRO-2: плюс градиенты | Обучение на нескольких ускорителях, когда обмен градиентами уже стал узким местом | Растёт объём обмена между устройствами, чувствительность к топологии связей |
| ZeRO-3: плюс параметры | Модель целиком не помещается в память ускорителей, нужны длинный контекст и крупный пакет | Параметры подтягиваются на каждом шаге — выше задержка и требования к каналу обмена |
| Выгрузка состояний в память процессора | Память ускорителей в дефиците, а шаг чувствителен к задержкам | Ограничена пропускной способностью шины между процессором и ускорителями |
| Выгрузка состояний на диск | Модель не помещается даже с полным шардированием параметров | Заметная просадка скорости шага, нужен быстрый локальный накопитель, растёт износ |
Как выбираем на практике
Сначала фиксируем бюджет памяти ускорителей и требуемый размер пакета данных, затем снимаем замер потребления на одном шаге обучения. Если веса помещаются — останавливаемся на первом или втором режиме, если нет — идём к полному шардированию и лишь затем к выгрузке. Смена режима — это один параметр конфигурации, поэтому откат к предыдущему варианту занимает минуты, а не переписывание кода обучения.
Порядок внедрения: от аудита нагрузки до передачи
Прежде чем что-то менять, нужно понять, где именно обучение упирается: в объём памяти или в обмен данными между ускорителями. Поэтому работа начинается с замера текущего режима, а не с правки конфигурации наугад — иначе эффект от изменений не с чем сравнить.
Семь этапов
- Аудит обучения и профиль нагрузки. Снимаем профиль одного шага: сколько памяти занимают веса, состояния оптимизатора и активации, какую долю времени ускорители простаивают в ожидании обмена. На выходе — узкое место, названное числом, а не догадка.
- Схема размещения. Подбираем степень секционирования состояний модели (ZeRO) и режим offloading: что остаётся на ускорителе, что уходит в память CPU, что пересчитывается на ходу. Схему фиксируем до правок.
- Внедрение обучения. Меняем параметры по одному: размер микро-батча, шаг накопления градиента, перекрытие обмена с вычислениями, частоту контрольных точек. После каждого изменения — замер.
- Серия контрольных прогонов. Короткие прогоны на реальном наборе данных со сверкой кривой сходимости с эталонной. Цель — убедиться, что выигрыш в скорости не куплен деградацией качества.
- Стабилизация длинных прогонов. Включаем восстановление из контрольных точек, ограничиваем рост их числа, проверяем поведение при отказе одного процесса — обучение продолжается, а не падает целиком.
- Документирование. Схема размещения, значения параметров и обоснование каждого изменения. Без этого конфигурация снова станет «чёрным ящиком» после первой же правки.
- Передача в эксплуатацию. Прогоняем обучение вместе с вашей командой, разбираем типовые сбои и точки правки, оставляем регламент запуска и изменения параметров.
Что остаётся у команды на выходе
Передаём не только работающий запуск, но и всё, что нужно, чтобы команда воспроизводила и сопровождала его без нас. Конфигурация обучения оформлена как набор профилей под конкретный состав железа, а не как удачный подбор параметров, который держится в голове одного инженера.
- Рабочие профили запуска: число процессов, размер порции данных на устройство, шаг накопления градиентов, ступень разбиения состояний модели и режим вынесения их на память процессора или диск.
- Схема размещения модели и состояний: что живёт на ускорителях, что вынесено за их пределы, сколько занимает каждое состояние и при каком размере модели схема перестаёт укладываться в ресурсы.
- Регламент запуска: последовательность проверок перед стартом — доступность всех устройств, свободный объём памяти, целостность последней сохранённой точки, совпадение профиля с фактическим железом.
- Регламент перезапуска: с какой точки продолжается обучение, как проверяется сохранённое состояние и за какое время процесс возвращается в работу после сбоя.
- Инструкция для дежурной смены: какие записи в журнале читать, какие признаки деградации считать аварийными, что исправляется на месте и когда вызывается ответственный.
- Журнал изменений конфигурации: каждая правка с причиной и замером до и после — время шага, потребление памяти, динамика качества на контрольной выборке.
- Схема оповещений: набор метрик и порогов, по которым смена узнаёт о проблеме раньше, чем обучение встанет на сутки.
Папка с профилями, схемами и регламентами лежит рядом с кодом и версионируется вместе с ним. Приёмку проводим вместе с вашей командой: дежурный инженер самостоятельно прогоняет запуск и восстановление на тестовом контуре, а мы отвечаем на вопросы по каждому пункту.
Кейс: как сняли ограничение по памяти при обучении
К нам пришла команда, которой не хватало памяти для следующей версии модели. Обучение падало не на первом шаге, а через несколько часов: замеры показали, что состояния оптимизатора и градиенты занимали больше памяти, чем сами веса, а часть ускорителей простаивала в ожидании обмена.
Что замеряли
Сняли профиль по шагам: пиковое потребление памяти, утилизацию каждого ускорителя, долю времени на пересылку между картами и число прогонов, доживавших до конца. Это дало базу для сравнения — без неё любые изменения в режиме размещения выглядят как вера, а не как инженерное решение.
Что изменили и что получили
Перевели обучение на DeepSpeed: включили шардирование состояний оптимизатора, градиентов и параметров между картами (семейство ZeRO), вынесли часть состояний на процессор и заменили рост микро-батча на накопление градиента. Активации пересчитываем на обратном проходе, а не держим в памяти целиком — этот обмен вычислений на память здесь выгоднее.
Размер модели вырос почти втрое на том же числе ускорителей. Доля прогонов, доходивших до конца, поднялась примерно с трети до всех, сбои по памяти ушли как класс, а время одного шага выросло на 10–15% — это осознанная плата за предсказуемость. Контрольные точки сохраняются и восстанавливаются, откат на предыдущую конфигурацию размещения занимает минуты.
Чем ZeRO отличается от выгрузки состояний на диск?
Это разные механизмы, а не альтернативы. ZeRO делит состояния модели — оптимизатор, градиенты, параметры — между ускорителями, и каждое устройство держит только свою долю. Выгрузка ничего не делит: она переносит состояния в память процессора или на диск, освобождая видеопамять.
Стадии ZeRO наращивают эффект: первая делит состояние оптимизатора, вторая добавляет градиенты, третья — сами параметры. Плата — обмен между ускорителями на каждом шаге: чем выше стадия, тем больше трафика и тем сильнее обучение зависит от пропускной способности межсоединения.
| Вариант | Что снимает с ускорителя | Цена | Когда подходит |
|---|---|---|---|
| ZeRO-1 | состояние оптимизатора | обмен между ускорителями | модель влезает, но оптимизатор тяжёлый |
| ZeRO-2 | оптимизатор и градиенты | трафик заметно растёт | нужен запас под крупные пакеты |
| ZeRO-3 | плюс сами параметры | максимальный трафик | иначе модель не помещается на группу |
| Выгрузка на CPU | оптимизатор, градиенты, буферы | пропускная способность шины | памяти мало, время шага терпит |
| Выгрузка на диск | состояния и параметры | самый медленный шаг | модель сильно превышает суммарную память |
Выгрузка размещает те же тензоры в памяти процессора или на диске и подкачивает их к моменту вычислений. Пиковое потребление видеопамяти падает, но шаг упирается в скорость передачи по шине — выигрыш в памяти оплачивается временем шага.
На практике механизмы комбинируют: в DeepSpeed они включаются независимо, и связка ZeRO-3 с выгрузкой параметров обучает модель, которая иначе не помещается на группу ускорителей. Мы считаем бюджет памяти под состояние и активации, затем подбираем режим и проверяем его на замерах реального шага.
Частые технические сомнения перед стартом
До старта инженерные команды проверяют три вещи: как настройка ляжет на существующий контур обучения, что будет с собранными чекпоинтами и как откатиться, если результат не устроит. Ниже — ответы на то, что чаще всего обсуждаем на первых созвонах.
Совместимо ли это с нашим контуром обучения?
Конфигурацию собираем вокруг вашего запуска: тренер, формат датасета и схема чекпоинтов остаются прежними. Меняется профиль запуска — глубина ZeRO, выгрузка состояний оптимизатора на CPU и диск, накопление градиентов, — а не код модели. Перед масштабированием сверяем один шаг обучения с текущим запуском.
Вырастет ли время обучения?
Выгрузка меняет память на скорость обмена, поэтому профиль подбираем по замерам: сколько шагов в секунду даёт каждый вариант и какой размер пакета помещается. Обычно освобождённая память позволяет поднять размер пакета и число устройств и перекрыть накладные расходы, но без замеров это не утверждение.
Как откатить изменения, если результат не устроит?
Изменения изолированы: конфигурация и обёртки запуска лежат отдельно, прежний запуск остаётся рабочим. Возврат — переключение на старый профиль и продолжение с последнего чекпоинта; формат чекпоинтов не меняем, поэтому промежуточные результаты читаются и прежним, и новым запуском.
Как убедиться, что качество не упало?
Сравниваем на коротком прогоне: те же данные и порядок, эталонный запуск и новый профиль, кривые потерь и метрики на отложенной выборке. Расхождение в пределах шума — продолжаем; систематический сдвиг — ищем причину в точности накопления градиентов и настройках оптимизатора.
Обсудим задачу и вернёмся с оценкой и планом
План и оценка появляются не «на глаз», а из входных данных. Чтобы посчитать этапы оптимизации обучения модели, нам нужно увидеть текущую конфигурацию запуска, профиль нагрузки на несколько ускорителей и требования к времени шага — без этого оценка останется вилкой, а не планом.
Для старта достаточно:
- доступ к репозиторию с текущими скриптами запуска и конфигами шардинга и выгрузки — достаточно прав на чтение;
- описание модели и её размера, число ускорителей и тип связи между ними;
- профиль реальной нагрузки: длина последовательностей, доля дообучения, требования к пропускной способности;
- журналы неудачных запусков или отчёты о нехватке памяти, если они сохранились;
- ограничения по средам: своя инфраструктура или облако, требования к доступам и изоляции;
- контакт инженера, который принимает результат, — чтобы согласовывать изменения без посредников.
Оценка возвращается документом: перечень этапов, ожидаемый эффект каждого (утилизация ускорителей, объём выгружаемых состояний, время шага), риски и точки отката. Первый этап начинается не с правки конфигурации, а с воспроизводимого замера: фиксируем базовую пропускную способность, снимаем профиль памяти, и только затем меняем стратегию шардинга и порядок выгрузки.
Опишите задачу в свободной форме — вернёмся с оценкой и планом работ.
