PyTorch FSDP: как ускорить обучение больших моделей и сэкономить память GPU

PyTorch Fully Sharded Data Parallel (FSDP) — это метод распределённого обучения, который шардирует параметры, градиенты и состояния оптимизатора между несколькими GPU. Благодаря этому можно обучать модели с миллиардами параметров даже при ограниченном объёме видеопамяти. FSDP является частью экосист

Как работает FSDP

FSDP реализует стратегию полного шардирования данных. В отличие от традиционного Distributed Data Parallel (DDP), где каждый GPU хранит полную копию модели, FSDP распределяет параметры, градиенты и состояния оптимизатора по всем доступным устройствам. Это кардинально снижает потребление памяти на каждом GPU, позволяя обучать значительно более крупные модели.

Процесс обучения с FSDP состоит из трёх этапов. На этапе forward каждый GPU собирает необходимые параметры из других устройств, выполняет прямой проход и затем освобождает ненужные параметры. На этапе backward градиенты собираются и усредняются по всем GPU, после чего каждый GPU обновляет свою часть параметров. На этапе update состояния оптимизатора также шардируются, что дополнительно экономит память.

Почему FSDP — это прорыв

С ростом размеров языковых моделей, таких как LLaMA и GPT, обучение на одном GPU стало практически невозможным. Например, модель с 10 миллиардами параметров требует около 40 ГБ памяти только для хранения весов в формате float32. С учётом градиентов и состояний оптимизатора потребление памяти легко превышает 80 ГБ, что недоступно даже для топовых GPU. FSDP решает эту проблему, распределяя нагрузку между несколькими устройствами.

Эксперименты показывают, что FSDP позволяет обучать модели с более чем 10 миллиардами параметров на кластере из 8 GPU A100 с 80 ГБ памяти каждая. Для сравнения, DDP на такой же конфигурации может обрабатывать модели лишь до 2–3 миллиардов параметров. При этом FSDP сохраняет высокую скорость обучения за счёт эффективного параллелизма.

Как настроить FSDP для вашей модели

Настройка FSDP в PyTorch относительно проста. Достаточно обернуть модель в специальный класс и указать политику шардирования. Например, для модели BLOOM-7B можно использовать следующий подход: установить shardingstrategy в FULLSHARD, что обеспечит полное шардирование всех параметров, градиентов и состояний оптимизатора. Для моделей с 13 миллиардами параметров, таких как LLaMA-13B, рекомендуется использовать гибридную стратегию, где часть данных остаётся локальной для ускорения обмена.

Важно правильно выбрать количество GPU и размер батча. FSDP эффективен при использовании от 4 до 64 GPU. При меньшем количестве устройств выигрыш в памяти может быть незначительным, а при большем — возрастают накладные расходы на коммуникацию. Рекомендуется начинать с 8 GPU и постепенно масштабировать.

Какие модели можно обучать с помощью FSDP

FSDP подходит для любых моделей на основе трансформеров, включая BLOOM, LLaMA, GPT-NeoX и другие. Особенно эффективен он для моделей с числом параметров от 1 миллиарда и выше. Для небольших моделей (менее 500 миллионов параметров) использование FSDP может быть избыточным, так как накладные расходы на шардирование превышают выгоду.

В официальном руководстве Hugging Face приведены примеры конфигураций для BLOOM-7B и LLaMA-13B. Для BLOOM-7B на 8 GPU A100 80GB удаётся достичь скорости обучения около 1000 токенов в секунду на GPU при размере батча 4. Для LLaMA-13B на той же конфигурации скорость составляет примерно 800 токенов в секунду на GPU.

FSDP vs DeepSpeed ZeRO: что выбрать?

DeepSpeed ZeRO — ещё один популярный метод шардирования, разработанный Microsoft. Оба подхода решают схожие задачи, но имеют различия. FSDP является встроенным решением PyTorch, что упрощает интеграцию и не требует установки дополнительных библиотек. DeepSpeed ZeRO предлагает более гибкие уровни оптимизации (ZeRO-1, ZeRO-2, ZeRO-3) и может быть быстрее на некоторых конфигурациях за счёт оптимизированных коммуникационных протоколов.

Однако точные сравнения производительности на одинаковых конфигурациях пока не опубликованы. Выбор между FSDP и DeepSpeed ZeRO часто зависит от конкретной задачи и предпочтений команды. Если вы уже используете PyTorch и хотите минимизировать зависимости, FSDP — отличный выбор. Если же нужна максимальная производительность и гибкость, стоит рассмотреть DeepSpeed ZeRO.

Какие ограничения есть у FSDP

На данный момент FSDP оптимизирован в первую очередь для GPU NVIDIA. Поддержка других аппаратных платформ, таких как AMD или Intel, ограничена. Кроме того, FSDP требует высокой пропускной способности межсоединений GPU (например, NVLink), иначе узким местом становится обмен данными. В облачных средах с медленной сетью эффективность FSDP может снижаться.

Ещё одно ограничение — сложность отладки. Из-за шардирования трассировка ошибок становится менее прозрачной. Рекомендуется сначала запускать обучение на небольшой модели для проверки корректности конфигурации.

Как начать использовать FSDP прямо сейчас

Чтобы начать работу с FSDP, установите последнюю версию PyTorch (1.12 или новее) и воспользуйтесь примером из официальной документации. Оберните модель в FullyShardedDataParallel, укажите количество GPU и запустите обучение. Для моделей Hugging Face можно использовать интеграцию через Trainer, который поддерживает FSDP из коробки.

Важно помнить, что FSDP требует большего объёма оперативной памяти CPU для хранения шардированных данных. Убедитесь, что на узле достаточно RAM (рекомендуется от 256 ГБ для моделей размером 10 миллиардов параметров).

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