Entraînement multi-GPU

LeRobot s’entraîne sur plusieurs GPU via Hugging Face Accelerate. Trois configurations de parallélisme de données sont prises en charge :

Installation

accelerate est inclus dans l’extra training :

pip install 'lerobot[training]'

Lancement

La formation distribuée peut être lancée à la fois via torchrun et accelerate launch. Accelerate est utilisé comme simple lanceur : il ne gère pas la configuration de l’entraînement, et tous les paramètres de formation distribuée sont gérés par le propre système de configuration de LeRobot.

Avec 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

Avec accelerate launch (en tant que lanceur simple) :

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

Sans les drapeaux --parallelism.*, un lancement multi-processus exécute un DDP standard. Les exécutions multi-nœuds utilisent les drapeaux standard torchrun --nnodes/--node-rank/--rdzv-endpoint (ou accelerate launch --num_machines/--machine_rank/--main_process_ip).

Les fichiers de configuration YAML d’Accelerate (accelerate launch --config_file some.yaml, accelerate config) ne sont pas pris en charge. Ils configurent le moteur via des variables d’environnement, en contournant le système de configuration de LeRobot, de sorte que train_config.json ne décrirait plus les paramètres réellement utilisés par une exécution. lerobot-train refuse donc de démarrer lorsque les variables d’environnement accelerate sont définies. Placez plutôt les paramètres dans les drapeaux --parallelism.* / --accelerator.*, ou définissez LEROBOT_ALLOW_ACCELERATE_ENV=1 pour prendre en compte la surcharge et continuer quand même.

Sémantique des lots, taux d'apprentissage et nombre d'étapes

Chacun des dp_replicate × dp_shard travailleurs parallèles de données charge son propre micro-lot --batch_size à chaque étape, de sorte qu'une étape d'entraînement consomme batch_size × dp_world_size échantillons, et × gradient_accumulation_steps de ceux-ci sont intégrés à chaque mise à jour de l'optimiseur :

effective_batch_size = batch_size × dp_world_size × gradient_accumulation_steps

La bannière d’entraînement affiche cette factorisation au démarrage. --steps compte les étapes de la boucle (micro-lots par travailleur), et non les mises à jour de l’optimiseur.

L'accumulation de gradients est un drapeau de premier ordre :

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

LeRobot ne met pas à l’échelle automatiquement le taux d’apprentissage ni le nombre d’étapes lorsque la taille effective du lot augmente. Si vous passez à une configuration distribuée et souhaitez un entraînement équivalent, veuillez effectuer les ajustements manuellement, par exemple avec 2 GPU : doublez --optimizer.lr (mise à l’échelle linéaire) ou divisez par deux --steps.

Entraînement fractionné (FSDP)

Si un modèle est trop volumineux pour être entraîné avec DDP, fragmentez-le avec 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 effectue le fractionnement sur le nombre de processus démarrés par le lanceur.

--accelerator.mixed_precision prend en charge à la fois bf16 et fp16. Détails dans mixed precision

Envelopper les unités

FSDP découpe le modèle en unités (généralement le bloc transformer répété) et rassemble une unité à la fois pendant les phases forward/backward. Les politiques déclarent leurs unités d’encapsulation via _fsdp_wrap_modules sur la classe de politique. Par exemple, ACT déclare ["ACTEncoderLayer", "ACTDecoderLayer"] et FastWAM déclare ["MoTLayer"]. Pour une politique sans déclaration _fsdp_wrap_modules, passez l’un des drapeaux ci-dessous. Vous pouvez spécifier explicitement le nom de la classe de module, ou utiliser à la place une politique basée sur la taille :

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

FSDP ne peut pas répartir les poids liés (par exemple, une tête de langage liée à l’embedding d’entrée) entre les unités d’emballage. torch 2.12 et les versions ultérieures s’arrêtent à la première étape d’entraînement et nomment le paramètre concerné, tandis que les versions antérieures les délient silencieusement. L’emballage de la classe de bloc répété les maintient généralement ensemble dans l’unité racine.

Si une politique ne déclare pas _fsdp_wrap_modules et qu’aucun argument --accelerator.fsdp.wrap_modules ou --accelerator.fsdp.min_num_params n’est passé, l’exécution échoue au démarrage plutôt que d’appliquer silencieusement l’encapsulation uniquement au module racine (ce qui ferait perdre tous les gains de mémoire liés au fractionnement).

Autres paramètres de sharding :

• --accelerator.fsdp.reshard_after_forward : indique s'il faut conserver les paramètres de chaque unité en mémoire après le passage avant.
• --accelerator.fsdp.cpu_offload : conserve les paramètres, les gradients et les états de l'optimiseur sur le CPU.
• --accelerator.fsdp.ignored_modules : une expression régulière de chemins de modules à conserver non fragmentés.
• --accelerator.fsdp.mixed_precision_policy.{param_dtype,reduce_dtype,output_dtype,cast_forward_inputs} : les champs de MixedPrecisionPolicy de PyTorch, appliqués à chaque unité FSDP. Dès qu’un champ est défini, la politique est construite uniquement à partir de ces champs et remplace celle dérivée de --accelerator.mixed_precision, il faut donc également définir param_dtype. Avant torch 2.15, un reduce_dtype non défini suit param_dtype. À partir de torch 2.15, il suit le dtype propre des paramètres (float32 en précision mixte). float16 nécessite --accelerator.mixed_precision=fp16, qui construit le loss scaler.

HSDP

Parallélisme de données hybride et fragmenté : les paramètres, les gradients et les états de l'optimiseur sont fragmentés sur dp_shard rangs, et cette fragmentation est répliquée dp_replicate fois. Les all-gathers de paramètres et les reduce-scatters de gradients restent à l'intérieur d'un groupe de fragments ; seul l'all-reduce qui synchronise les répliques traverse les groupes. Les deux degrés doivent multiplier la taille du monde :

# 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

La fine-tuning PEFT (--peft.*) fonctionne sous FSDP2 et HSDP. Le modèle de base gelé est sharded avec les adaptateurs, de sorte que chaque rang ne détient que son shard des poids de base :

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

• Les unités enveloppées proviennent toujours de la déclaration _fsdp_wrap_modules de la politique (ou de --accelerator.fsdp.wrap_modules).
• Utilisez la précision mixte lorsque le modèle de base est stocké en bf16 ou fp16. PEFT conserve ses adaptateurs en fp32, tandis que les copies de --peft.full_training_modules conservent le dtype du modèle de base, et FSDP2 a besoin d'un dtype unique pour les paramètres entraînables qu'il sharde ensemble. Avec la précision mixte activée, accelerate convertit chaque paramètre entraînable en fp32.
• Les points de contrôle stockent l'adaptateur plutôt que les poids complets, dans chaque --checkpoint_format (voir Points de contrôle).
• --accelerator.fsdp.ignored_modules est comparé aux noms de modules du modèle PEFT, qui commencent par base_model.model.. Commencez le motif par .* pour les correspondre.

Points de contrôle

Chaque point de contrôle contient un répertoire pretrained_model/ et un répertoire 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

Lors de l’entraînement sur un seul GPU ou avec DDP, le pipeline sérialise chaque state dict dans un fichier unique : model.safetensors pour le modèle et optimizer_state.safetensors pour l’optimiseur.

Lors de l’entraînement sharded, l’état de l’optimiseur est enregistré sous forme de shards DCP sous training_state/optimizer_0/, et la disposition du modèle sous pretrained_model/ peut être configurée via --checkpoint_format :

Une exécution PEFT stocke son adaptateur au lieu des poids complets : adapter_model.safetensors et adapter_config.json remplacent model.safetensors, et les fragments DCP ne contiennent que l’adaptateur. Le modèle de base est chargé depuis le chemin enregistré dans adapter_config.json.

Deux points à connaître concernant les points de contrôle regroupés (safetensors) issus d'exécutions fractionnées :

• Ils stockent les poids en fp32. Lors d’une formation en précision mixte, FSDP conserve une copie maîtresse en fp32, et le point de contrôle enregistre cette copie maîtresse afin de garantir que la formation reprend de manière cohérente.
• La collecte est collective (tous les rangs participent), mais seul le processus principal écrit.

Conversion des points de contrôle DCP

lerobot-convert-dcp fusionne un répertoire de shards DCP en un model.safetensors standard, hors ligne et sans 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 publie le répertoire converti en tant que dépôt de modèle. Les shards d'une exécution PEFT fusionnent dans adapter_model.safetensors. Comme tout checkpoint dcp, un checkpoint PEFT enregistré avec --checkpoint_format=dcp se charge avec --policy.path uniquement après cette conversion.

Reprise

Reprendre avec --resume=true --config_path=.../checkpoints/last/pretrained_model/train_config.json. La reprise à partir d'un point de contrôle DCP prend en charge le re-morcelage (resharding) du modèle et de l'état de l'optimiseur vers la topologie actuelle, ce qui signifie que vous pouvez reprendre avec une répartition dp_replicate/dp_shard différente. L'échantillonneur de données peut toujours reprendre à la bonne époque et au bon décalage, mais il n'est exact à l'échantillon près que si la taille du monde et la taille du lot correspondent à celles de l'exécution d'origine (un avertissement est journalisé dans le cas contraire).

Les points de contrôle FSDP écrits par LeRobot 0.6.x et les versions antérieures utilisaient une structure sur disque différente (un état d'optimiseur complet regroupé) et ne peuvent pas être repris.

Notes

• Les sauvegardes de points de contrôle et les publications en fin d’entraînement sont collectives (chaque rang y entre). Les poids agrégés, les fichiers secondaires et les téléversements sur le Hub sont écrits uniquement par le processus principal.
• Les métriques sont réduites entre les rangs avant la journalisation : les pertes sont moyennées et samples/s rapporte le débit de l'ensemble du cluster.
• Le planificateur de taux d'apprentissage est incrémenté une fois par étape d'entraînement, indépendamment du nombre de processus (step_scheduler_with_optimizer=False est intégré).

Pour un contexte sur le mécanisme sous-jacent, consultez le guide Accelerate FSDP. Pour approfondir l’entraînement à grande échelle, consultez le Ultrascale Playbook.