Распределённое сохранение состояния PyTorch (DCP)
Определение
DCP — средство PyTorch, с помощью которого несколько процессов сохраняют и загружают переданное состояние модели и оптимизатора. При загрузке части тензоров можно перераспределить между процессами; целевые тензоры подготавливают заранее.
Также ищут: DCP · Distributed Checkpoint · PyTorch Distributed Checkpoint · torch.distributed.checkpoint · Распределённые контрольные точки PyTorch
Зачем это знать
DCP позволяет загрузить совместимое состояние при другом числе процессов. Сверьте ключи в их словарях, создайте тензоры на целевых устройствах и включите в контрольную точку нужные данные приложения. Успешная загрузка сама по себе не подтверждает совпадение дальнейшего обучения.
Карта понятия
Что подтверждено
PyTorch DCP параллельно сохраняет распределённое состояние: каждый участник записывает локальные фрагменты ShardedTensor и DTensor. При загрузке DCP перераспределяет фрагменты, что позволяет сохранить состояние при одной топологии кластера и загрузить при другой.
Граница применимости: torch.distributed.checkpoint в официальной документации PyTorch 2.9. Речь о локальных частях поддерживаемых ShardedTensor и DTensor и загрузке совместимого состояния; это не обещание переноса произвольной архитектуры модели или одинакового пика памяти. Равное деление в учебном примере задано отдельно, без реплик.
В PyTorch v2.9.1 get_state_dict использует канонические полные имена (FQN) параметров и буферов модели, а идентификаторы параметров в состоянии оптимизатора заменяет такими именами; для полученных состояний заявлено перераспределение при другом числе участников и схемах параллелизма. Парная функция set_state_dict применяет словари с такими именами к модели и оптимизатору. Это зависит от корректного состояния той же модели; канонические имена не делают произвольное изменение архитектуры совместимым.
Граница применимости: Реализация torch.distributed.checkpoint.state_dict.py в теге PyTorch v2.9.1. Описаны модель и связанный оптимизатор; поддержка смены числа участников не означает автоматического преобразования несовместимых ключей или архитектур.
Перед вызовом DCP load все участники должны передать state_dict с одинаковыми ключами, а тензоры в нём должны быть заранее выделены на целевых устройствах. Несовпадение ключей может привести к зависанию или ошибке.
Граница применимости: Вызов torch.distributed.checkpoint.load из официальной документации PyTorch 2.9. Предусловия относятся к переданному state_dict и распределённой загрузке; документация не обещает одинаковую причину отказа для каждого нарушения.
DCP сохраняет переданный state_dict. При сохранении DCP вызывает state_dict переданного объекта Stateful. При загрузке он сначала вызывает state_dict, затем заполняет полученный словарь и вызывает load_state_dict. Для обычных словарей с каноническими именами приложение может применить загруженное состояние модели и оптимизатора через set_state_dict. Планировщик обучения, позиция данных, генераторы случайных чисел и другие данные попадут в контрольную точку, только если приложение включит их в передаваемое состояние. PyTorch не гарантирует обратную совместимость сохранённых state_dict между версиями.
Граница применимости: Сохранение и загрузка через torch.distributed.checkpoint в документации PyTorch 2.9; роль set_state_dict дополнительно подтверждена реализацией API PyTorch v2.9.1. Утверждение о данных приложения — граница явно переданного state_dict, а не полный перечень API или обещание точного продолжения обучения.
Что проверить
- Сверьте ключи state_dict у всех процессов, участвующих в загрузке; несовпадение может вызвать зависание или ошибку.
- Создайте целевые тензоры на нужных устройствах до вызова load.
- Для словарей с каноническими именами проверьте применение загруженного состояния через set_state_dict. Для Stateful проверьте метод load_state_dict: DCP вызывает его после чтения, а обёртка применяет полученные значения.
- Включите нужные дополнительные данные приложения в сохраняемое состояние. Проверьте загрузку на целевой версии PyTorch и последующие шаги обучения.
Первоисточники
- Distributed Checkpoint — torch.distributed.checkpoint
Официальная документация PyTorch 2.9: параллельное сохранение локальных фрагментов, переразложение при загрузке, входное состояние, условия загрузки и отсутствие гарантии обратной совместимости сохранённых state_dict между версиями PyTorch.
- torch.distributed.checkpoint.state_dict
Официальный исходный код из тега PyTorch v2.9.1: get_state_dict и set_state_dict возвращают и применяют состояния модели и оптимизатора, согласованные по каноническим FQN; описана поддержка перераспределения при другом числе участников и сочетании схем параллелизма.
Связанные понятия
Курсы по теме
В каталоге
- GPU для продолжения обучения — Сопоставьте память ускорителей с новым распределением состояния и измеренным пиком памяти шага после загрузки. Сам перенос состояния не подтверждает достаточность памяти.