Checkpoint sharding — раздельное сохранение весов модели на несколько файлов или устройств, уменьшающее требования к памяти и ускоряющее загрузку.
Определение
Checkpoint sharding — это метод сохранения модели, при котором веса, градиенты и состояния оптимизатора разбиваются на части (шарды) и записываются в разные файлы или на разные устройства. Такой подход используется при работе с моделями, объём которых превышает возможности локальной памяти, а также в распределённых тренинговых системах, где параметры изначально хранятся не целиком на одном GPU.
В современных системах checkpoint sharding стал стандартным способом работы с моделями объёмом десятки и сотни гигабайт, для которых обычное сохранение в один файл невозможно или неэффективно.
Как работает
При сохранении состояния обучения фреймворк делит параметры на части с учётом распределения модели между устройствами. Каждый GPU сохраняет свой фрагмент весов, градиентов или оптимизаторных состояний. В дальнейшем при загрузке шарды объединяются или загружаются по запросу.
Основные механизмы:
- Разбиение параметров — каждый слой или блок делится на несколько файлов.
- Сохранение локальных шардов — GPU записывает только свою часть параметров.
- Индексы шардов — метаданные связывают файлы между собой и определяют структуру модели.
- On-demand загрузка — при восстановлении обучающая система загружает только нужные фрагменты.
- Оптимизация ввода-вывода — шардинг снижает нагрузку на файловую систему.
Checkpoint sharding используется как при model/tensor parallelism, так и при ZeRO-оптимизациях, где часть параметров вообще отсутствует на большинстве GPU во время тренировки.
Где применяется
- Обучение LLM с параметрами, превышающими локальную память GPU.
- DeepSpeed ZeRO-2 и ZeRO-3.
- PyTorch FSDP и Megatron-LM.
- Сохранение моделей для многоузловых тренингов.
- Хранение огромных чекпоинтов (50–500+ ГБ) в распределённых средах.
- Деплой больших моделей на кластерах с ограниченным I/O.
Практические примеры использования
В PyTorch FSDP каждый GPU хранит свой шард параметров слоя. При сохранении checkpoint фреймворк записывает множество файлов вида shard_0.pt, shard_1.pt и индекс, который описывает структуру. Это позволяет быстро сохранять и восстанавливать модели с сотнями миллиардов параметров.
В DeepSpeed ZeRO-3 каждый GPU хранит только часть оптимизаторных состояний и параметров. Чекпоинт формируется из большого набора шардов, содержащих разбитые тензоры и метаданные для сборки.
В Megatron-LM checkpoint sharding применяется вместе с tensor parallelism: матрицы, разделённые по колонкам или строкам, сохраняются отдельными файлами, соответствующими GPU-ранкам.
В inference-системах, где модель слишком велика для одного сервера, checkpoint sharding позволяет загружать и использовать веса частями, распределяя их между устройствами.
Ключевые свойства checkpoint sharding
- Линейная масштабируемость — большие модели сохраняются через распределённую систему.
- Эффективность I/O — уменьшение нагрузки на файловую систему благодаря параллельной записи.
- Совместимость с параллелизмом — подходит для tensor, pipeline, ZeRO и FSDP.
- Гибкая загрузка — возможно восстановление отдельных фрагментов.
Проблемы и ограничения
- Сложность организации — требуется корректная работа индексов и метаданных.
- Большое количество файлов — для крупных моделей сотни или тысячи шардов.
- Зависимость от стабильности файловой системы — потеря шарда может нарушить восстановление.
- Требования к сети — для восстановления на других узлах нужна быстрая передача файлов.
- Неравномерная загрузка — неверная стратегия шардинга может создавать дисбаланс.
Преимущества и ограничения
- Плюс: позволяет сохранять модели, которые невозможно записать в один файл.
- Плюс: снижает объем потребляемой памяти на GPU.
- Плюс: ускоряет сохранение больших чекпоинтов за счёт параллельной записи.
- Плюс: оптимально работает совместно с ZeRO и FSDP.
- Минус: требует сложной системы индексации.
- Минус: увеличивает количество файлов.
- Минус: замедляет восстановление при слабом I/O.
- Минус: требует согласованности устройств и шардов.
Связанные термины
- ZeRO optimization
- FSDP (Fully Sharded Data Parallel)
- Tensor parallelism
- Pipeline parallelism
- Sharded training
- Distributed training
- Checkpointing
- All-gather