멀티-GPU 학습

LeRobot는 Hugging Face Accelerate를 통해 여러 GPU에서 학습합니다. 세 가지 데이터 병렬 레이아웃을 지원합니다:

설치

accelerate는 training extra에 포함되어 있습니다:

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)를 사용합니다.

Accelerate의 YAML 설정 파일(accelerate launch --config_file some.yaml, accelerate config)은 지원되지 않습니다. 이러한 파일은 환경 변수를 통해 엔진을 구성하여 LeRobot의 설정 시스템을 우회하므로, train_config.json이 실제 실행에 사용된 설정을 더 이상 설명하지 못하게 됩니다. 따라서 accelerate 환경 변수가 설정된 경우 lerobot-train은 시작을 거부합니다. 대신 --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는 옵티마이저 업데이트가 아닌 루프 단계(워커당 마이크로 배치)를 세웁니다.

Gradient accumulation은 1급 플래그입니다:

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

LeRobot는 유효 배치 크기가 증가할 때 학습률이나 스텝 수를 자동으로 스케일링하지 않습니다. 스케일아웃하여 동등한 학습을 원할 경우, 예를 들어 GPU 2개를 사용할 때 --optimizer.lr을 두 배로 늘리거나(선형 스케일링) --steps를 절반으로 줄이는 방식으로 수동으로 조정해 주세요.

Sharded training (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를 모두 지원합니다. 자세한 내용은 혼합 정밀도를 참고하세요.

Wrap units

FSDP는 모델이 반복되는 transformer 블록과 같은 단위별로 모델을 분할(shard)하고, forward/backward 과정에서 한 번에 하나의 단위를 수집(gather)합니다. 정책(policy) 클래스의 _fsdp_wrap_modules를 통해 wrap 단위를 선언합니다. 예를 들어, 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: forward 이후 각 유닛의 파라미터를 유지할지 여부.
• --accelerator.fsdp.cpu_offload: 매개변수, 기울기 및 옵티마이저 상태를 CPU에 유지합니다.
• --accelerator.fsdp.ignored_modules: 샤딩하지 않고 유지할 모듈 경로에 대한 정규식입니다.
• --accelerator.fsdp.mixed_precision_policy.{param_dtype,reduce_dtype,output_dtype,cast_forward_inputs}: PyTorch의 MixedPrecisionPolicy 필드로, 모든 FSDP 유닛에 적용됩니다. 어떤 필드라도 설정되면, 해당 필드만으로 정책이 구성되며 --accelerator.mixed_precision에서 파생된 정책을 대체하므로 param_dtype도 함께 설정해야 합니다. torch 2.15 이전에는 reduce_dtype가 설정되지 않으면 param_dtype를 따릅니다. torch 2.15부터는 매개변수 자체의 dtype(float32 mixed precision 환경)을 따릅니다. float16은 loss scaler를 구성하는 --accelerator.mixed_precision=fp16이 필요합니다.

HSDP

Hybrid Sharded Data Parallel: 파라미터, 그래디언트 및 옵티마이저 상태는 dp_shard 랭크에 걸쳐 샤딩되며, 이 샤딩은 dp_replicate 번 복제됩니다. 파라미터 올게더(all-gather)와 그래디언트 리듀스스캐터(reduce-scatter)는 샤드 그룹 내에서만 수행되며, 레플리카를 동기화하는 올리듀스(all-reduce)만 그룹 간을 오갑니다. 두 차원의 곱이 월드 사이즈(world size)와 같아야 합니다:

# 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 fine-tuning(--peft.*)는 FSDP2와 HSDP에서 작동합니다. 동결된 base model은 adapter와 함께 sharding되며, 각 rank는 base weights의 자기 shard만 보유합니다:

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 unit은 여전히 정책의 _fsdp_wrap_modules 선언(또는 --accelerator.fsdp.wrap_modules)에서 가져옵니다.
• 기본 모델이 bf16 또는 fp16로 저장된 경우 혼합 정밀도를 사용하세요. PEFT는 어댑터를 fp32로 유지하는 반면, --peft.full_training_modules의 복사본은 기본 모델의 dtype을 유지합니다. FSDP2는 함께 샤딩하는 학습 가능 매개변수에 대해 단일 dtype이 필요합니다. 혼합 정밀도를 켜면 accelerate는 모든 학습 가능 매개변수를 fp32로 업캐스트합니다.
• --checkpoint_format의 모든 경우에서 체크포인트는 전체 가중치가 아닌 어댑터를 저장합니다(Checkpoints 참고).
• --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 학습 중에 파이프라인은 각 상태 딕셔너리를 단일 파일로 직렬화합니다: 모델은 model.safetensors, 옵티마이저는 optimizer_state.safetensors입니다.

sharded 학습 중에는 옵티마이저 상태가 training_state/optimizer_0/ 아래에 DCP 샤드로 저장되며, pretrained_model/ 아래 모델 레이아웃은 --checkpoint_format를 통해 구성할 수 있습니다:

PEFT 실행은 전체 가중치가 아닌 어댑터를 저장합니다: adapter_model.safetensors와 adapter_config.json이 model.safetensors를 대체하며, DCP 샤드는 어댑터만 포함합니다. 베이스 모델은 adapter_config.json에 기록된 경로에서 로드됩니다.

분할 실행에서 생성된 수집된 (safetensors) 체크포인트에 대해 알아야 할 두 가지 사항:

• fp32 가중치를 저장합니다. 혼합 정밀도 학습에서 FSDP는 fp32 마스터 복사본을 유지하며, 체크포인트는 훈련이 일관되게 재개되도록 마스터 복사본을 저장합니다.
• gather는 집합 연산(모든 rank가 참여)이지만, 메인 프로세스만 데이터를 기록합니다.

DCP 체크포인트 변환

lerobot-convert-dcp는 DCP 샤드 디렉터리를 GPU 없이 오프라인에서 일반적인 model.safetensors 파일로 병합합니다:

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는 변환된 디렉터리를 모델 저장소로 게시합니다. PEFT 실행의 샤드는 adapter_model.safetensors로 병합됩니다. 모든 dcp 체크포인트와 마찬가지로, --checkpoint_format=dcp로 저장된 PEFT 체크포인트는 이 변환 후에만 --policy.path로 로드됩니다.

재개

--resume=true --config_path=.../checkpoints/last/pretrained_model/train_config.json을 사용하여 재개합니다. DCP 체크포인트에서 재개하는 경우 모델과 옵티마이저 상태를 현재 토폴로지에 맞게 리샤딩하는 것을 지원하므로, 다른 dp_replicate/dp_shard 분할로 재개할 수 있습니다. 데이터 샘플러는 항상 올바른 에포크와 오프셋에서 재개할 수 있지만, 월드 사이즈와 배치 사이즈가 원래 실행과 일치할 때만 샘플 단위 정확도를 보장합니다(일치하지 않으면 경고가 기록됩니다).

LeRobot 0.6.x 및 이전 버전에서 작성한 FSDP 체크포인트는 다른 디스크 레이아웃(집합된 전체 옵티마이저 상태)을 사용하므로 재개할 수 없습니다.

The input is a single word "Notes". 참고

• Checkpoint 저장과 학습 종료 시 게시는 집합 연산(모든 랭크가 진입)입니다. 수집된 가중치, 사이드카 파일 및 Hub 업로드는 메인 프로세스에서만 작성됩니다.
• 로그 기록을 위해 메트릭은 랭크 간에 축소됩니다. 손실은 평균화되고, samples/s는 클러스터 전체의 처리량을 보고합니다.
• 학습률 스케줄링은 프로세스 수와 관계없이 학습 단계마다 한 번씩 진행됩니다(step_scheduler_with_optimizer=False가 기본으로 적용됨).

기본 메커니즘에 대한 배경 지식을 보려면 Accelerate FSDP 가이드를 참고하세요. 대규모 학습에 대해 더 깊이 알고 싶다면 Ultrascale Playbook를 확인해 보세요.