Обучение на нескольких GPU

LeRobot обучается на нескольких GPU с помощью Hugging Face Accelerate. Поддерживаются три схемы параллелизма данных:

Установка

accelerate входит в дополнительный пакет training:

pip install 'lerobot[training]'

Запуск

Распределённое обучение можно запускать как через torchrun, так и через accelerate launch. Accelerate используется как обычный лаунчер: он не управляет конфигурацией обучения, а все настройки распределённого обучения хранятся в собственной системе конфигурации LeRobot.

С помощью torchrun:

torchrun --nproc-per-node=2 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=act \
  --policy.repo_id=${HF_USER}/my_trained_policy \
  --output_dir=outputs/train/act_multi_gpu \
  --job_name=act_multi_gpu \
  --trackio.enable=true

С помощью accelerate launch (как обычный запускатель):

accelerate launch --num_processes=2 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=act \
  --policy.repo_id=${HF_USER}/my_trained_policy \
  --output_dir=outputs/train/act_multi_gpu \
  --job_name=act_multi_gpu \
  --trackio.enable=true

Без флагов --parallelism.* запуск в многопроцессном режиме выполняет обычный DDP. Для запуска на нескольких узлах используются стандартные флаги torchrun --nnodes/--node-rank/--rdzv-endpoint (или accelerate launch --num_machines/--machine_rank/--main_process_ip).

Файлы конфигурации YAML для Accelerate (accelerate launch --config_file some.yaml, accelerate config) не поддерживаются. Они настраивают движок через переменные окружения, минуя систему конфигурации LeRobot, из-за чего train_config.json перестанет описывать настройки, фактически использованные в ходе запуска. Поэтому lerobot-train отказывается запускаться, если заданы переменные окружения accelerate. Вместо этого передавайте настройки через флаги --parallelism.* / --accelerator.* либо установите LEROBOT_ALLOW_ACCELERATE_ENV=1, чтобы подтвердить переопределение и продолжить работу.

Семантика батчей, скорость обучения и шаги

Каждый из dp_replicate × dp_shard рабочих процессов с параллелизмом данных загружает собственный микробатч --batch_size на каждом шаге, поэтому один шаг обучения потребляет batch_size × dp_world_size образцов, а × gradient_accumulation_steps из них попадают в каждое обновление оптимизатора:

effective_batch_size = batch_size × dp_world_size × gradient_accumulation_steps

Панель обучения выводит эту факторизацию при запуске. --steps считает шаги цикла (микробатчи на воркер), а не обновления оптимизатора.

Накопление градиентов — это флаг первого класса:

torchrun --nproc-per-node=2 $(which lerobot-train) \
  --batch_size=8 --accelerator.gradient_accumulation.steps=4 ...

LeRobot не масштабирует автоматически скорость обучения или количество шагов при увеличении эффективного размера батча. Если вы масштабируете систему и хотите сохранить эквивалентное обучение, отрегулируйте параметры вручную, например, при использовании 2 GPU: удвойте --optimizer.lr (линейное масштабирование) или уменьшите --steps вдвое.

Обучение с шардированием (FSDP)

Если модель слишком велика для обучения с DDP, разбейте её на шарды с помощью FSDP2:

torchrun --nproc-per-node=4 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=<your_policy> \
  --parallelism.dp_shard=4 \
  --accelerator.mixed_precision=bf16 \
  --output_dir=outputs/train/my_policy_fsdp

--parallelism.dp_shard=-1 распределяет данные по количеству процессов, запущенных лаунчером.

--accelerator.mixed_precision поддерживает как bf16, так и fp16. Подробности в разделе смешанная точность

Обёртка единиц

FSDP разбивает модель на блоки (как правило, повторяющиеся трансформерные блоки) и собирает по одному блоку за раз во время прямого и обратного проходов. Политики указывают свои блоки обёртки через _fsdp_wrap_modules в классе политики. Например, ACT объявляет ["ACTEncoderLayer", "ACTDecoderLayer"], а FastWAM — ["MoTLayer"]. Для политики без объявления _fsdp_wrap_modules передайте один из флагов ниже. Вы можете явно указать имя класса модуля или использовать политику на основе размера:

--accelerator.fsdp.wrap_modules='["MyTransformerBlock"]'   # explicit class names
--accelerator.fsdp.min_num_params=1000000                  # or: wrap every submodule above 1M params

FSDP не может распределять связанные веса (например, языковую голову, связанную со входным эмбеддингом) между единицами обёртки. В torch 2.12 и более новых версиях обучение останавливается на первом шаге с указанием имени параметра, тогда как в более старых версиях связь между весами молча снимается. Как правило, обёртка класса повторяющегося блока позволяет сохранить связанные веса в пределах корневого блока.

Если политика не объявляет _fsdp_wrap_modules и не переданы параметры --accelerator.fsdp.wrap_modules или --accelerator.fsdp.min_num_params, запуск завершается с ошибкой на этапе инициализации, а не молча оборачивает только корневой модуль (что лишило бы всех преимуществ экономии памяти за счёт шардирования).

Другие настройки шардирования:

• --accelerator.fsdp.reshard_after_forward: нужно ли удерживать параметры каждого блока в памяти после прямого прохода.
• --accelerator.fsdp.cpu_offload: сохраняет параметры, градиенты и состояния оптимизатора на CPU.
• --accelerator.fsdp.ignored_modules: регулярное выражение для путей модулей, которые нужно оставить без шардирования.
• --accelerator.fsdp.mixed_precision_policy.{param_dtype,reduce_dtype,output_dtype,cast_forward_inputs}: поля MixedPrecisionPolicy из PyTorch, применяемые к каждой единице FSDP. Как только задано хотя бы одно поле, политика строится исключительно на основе этих полей и заменяет ту, что выводится из --accelerator.mixed_precision, поэтому также установите param_dtype. До версии torch 2.15, если reduce_dtype не задано, оно следует за param_dtype. Начиная с torch 2.15, оно следует за собственным dtype параметров (float32 при смешанной точности). Для float16 требуется --accelerator.mixed_precision=fp16, который настраивает масштабатор потерь.

HSDP

Гибридный шардированный параллелизм данных: параметры, градиенты и состояния оптимизатора шардируются по рангам dp_shard, а это шардирование реплицируется dp_replicate раз. Все операции all-gather для параметров и reduce-scatter для градиентов выполняются внутри группы шардов; только all-reduce, синхронизирующий реплики, выполняется между группами. Произведение этих двух степеней должно быть равно общему размеру мира:

# 16 GPUs = 2 nodes × 8: shard within each node, replicate across nodes
torchrun --nnodes=2 --nproc-per-node=8 ... $(which lerobot-train) \
  --parallelism.dp_replicate=2 --parallelism.dp_shard=8 ...

PEFT

Дополнительная настройка PEFT (--peft.*) работает в FSDP2 и HSDP. Замороженная базовая модель разбивается на шарды вместе с адаптерами, поэтому каждый ранг хранит только свой шард базовых весов:

torchrun --nproc-per-node=4 $(which lerobot-train) \
  --policy.path=${HF_USER}/my_base_policy \
  --peft.method_type=LORA --peft.r=16 \
  --parallelism.dp_shard=4 \
  --accelerator.mixed_precision=bf16 \
  --output_dir=outputs/train/my_policy_fsdp_lora

• Wrap-модули по-прежнему определяются декларацией _fsdp_wrap_modules в политике (или параметром --accelerator.fsdp.wrap_modules).
• Используйте смешанную точность, если базовая модель хранится в bf16 или fp16. PEFT сохраняет свои адаптеры в fp32, тогда как копии --peft.full_training_modules сохраняют тип данных базовой модели, а FSDP2 требует единого типа данных для обучаемых параметров, которые он разделяет. При включенной смешанной точности accelerate приводит все обучаемые параметры к fp32.
• Чекпоинты хранят адаптер вместо полных весов в каждом --checkpoint_format (см. Чекпоинты).
• --accelerator.fsdp.ignored_modules сопоставляется с именами модулей модели PEFT, которые начинаются с base_model.model.. Начните шаблон с .*, чтобы они совпали.

Чекпоинты

Каждая контрольная точка содержит каталог pretrained_model/ и каталог training_state/:

005000/  # the training step at that checkpoint
├── pretrained_model/
│   ├── config.json  # policy config
│   ├── train_config.json  # the full training config
│   ├── model.safetensors  # full weights (checkpoint_format ∈ {safetensors, safetensors_dcp}, or any non-sharded run)
│   ├── pytorch_model_fsdp_0/  # DCP weight shards (checkpoint_format ∈ {dcp, safetensors_dcp})
│   ├── policy_preprocessor.json  # preprocessor config (when the run has a preprocessor)
│   ├── policy_preprocessor_step_*.safetensors  # state of the stateful preprocessor steps
│   ├── policy_postprocessor.json  # postprocessor config (when the run has a postprocessor)
│   └── policy_postprocessor_step_*.safetensors  # state of the stateful postprocessor steps
└── training_state/
    ├── optimizer_param_groups.json  # optimizer param groups (non-sharded runs)
    ├── optimizer_state.safetensors  # optimizer state (non-sharded runs)
    ├── optimizer_0/  # DCP optimizer shards (sharded runs)
    ├── rng_state.safetensors  # rng states
    ├── scheduler_state.json  # scheduler state (when the run has a scheduler)
    ├── scaler_state.json  # fp16 loss-scaler state (mixed_precision=fp16 only)
    └── training_step.json  # step + dp_world_size/batch_size/grad_accum/precision + topology

При обучении на одном GPU или с DDP конвейер сериализует каждый state dict в отдельный файл: model.safetensors для модели и optimizer_state.safetensors для оптимизатора.

Во время обучения с шардированием состояние оптимизатора сохраняется в виде DCP-шардов в training_state/optimizer_0/, а расположение модели в pretrained_model/ можно настроить с помощью --checkpoint_format:

Запуск PEFT сохраняет адаптер вместо полных весов: adapter_model.safetensors и adapter_config.json заменяют model.safetensors, а сегменты DCP содержат только адаптер. Базовая модель загружается из пути, указанного в adapter_config.json.

Два момента, которые нужно знать о собранных (safetensors) контрольных точках из разбитых на части запусков:

• Они хранят веса в формате fp32. При обучении с смешанной точностью FSDP сохраняет мастер-копию в fp32, а контрольная точка записывает эту мастер-копию, чтобы обеспечить согласованное возобновление обучения.
• Операция gather является коллективной (в ней участвуют все ранги), но запись выполняет только основной процесс.

Конвертация контрольных точек DCP

lerobot-convert-dcp объединяет каталог с DCP-фрагментами в обычный файл model.safetensors в автономном режиме и без использования GPU:

lerobot-convert-dcp --checkpoint_dir=outputs/train/run/checkpoints/005000
lerobot-convert-dcp --checkpoint_dir=... --delete_dcp=true --push_to_hub=${HF_USER}/my_policy

--push_to_hub публикует преобразованный каталог в виде репозитория модели. Шards запуска PEFT объединяются в adapter_model.safetensors. Как и любой checkpoint dcp, checkpoint PEFT, сохранённый с --checkpoint_format=dcp, загружается с --policy.path только после этого преобразования.

Возобновление

Возобновите с помощью --resume=true --config_path=.../checkpoints/last/pretrained_model/train_config.json. Возобновление из контрольной точки DCP поддерживает перераспределение модели и состояния оптимизатора под текущую топологию, что означает возможность возобновления с другим разбиением dp_replicate/dp_shard. Сэмплер данных всегда может возобновиться с правильной эпохой и смещением, но он является sample-exact только тогда, когда размер мира и размер батча совпадают с исходным запуском (в противном случае в журнал выводится предупреждение).

Чекпоинты FSDP, созданные в LeRobot 0.6.x и более ранних версиях, использовали другую структуру на диске (полное состояние оптимизатора, собранное воедино) и не могут быть восстановлены.

Примечания

• Сохранение контрольных точек и публикация в конце обучения являются коллективными (в них входят все ранги). Собранные веса, вспомогательные файлы и загрузка в Hub выполняются только основным процессом.
• Метрики агрегируются по рангам перед записью в журнал: потери усредняются, а samples/s отражает пропускную способность всего кластера.
• Планирование скорости обучения выполняется один раз на каждый шаг обучения, независимо от количества процессов (step_scheduler_with_optimizer=False встроено).

Для ознакомления с базовыми механизмами см. Accelerate FSDP guide. Чтобы углубиться в вопросы обучения в больших масштабах, обратитесь к Ultrascale Playbook.