Entraînement distribué dans NeMo AutoModel

Objectif

NeMo AutoModel utilise l'entraînement distribué natif de PyTorch. Tout le parallélisme est orchestré via un seul objet MeshContext qui contient les mailles de périphériques, les configurations de stratégie et les noms d'axes.

Instructions

Pour les questions conceptuelles sur l'entraînement distribué, répondez directement à partir des modèles rapides de cette skill sans inspecter le repository. Commencez par le choix de stratégie, puis énumérez uniquement les champs YAML et les contraintes pertinents à la question.

Utilisez des verbes d'action directs dans la réponse finale : recommandez la stratégie, montrez le YAML minimal, énoncez la contrainte de dimensionnement et nommez les stratégies non supportées. Ne discutez pas d'onboarding de modèle, de recipes, de Slurm, SkyPilot ou de checkpointing sauf si l'utilisateur le demande.

Exemples

TP plus PP pour un grand modèle multi-nœud

Recommandez strategy: fsdp2. Mentionnez tp_size, pp_size, cp_size, ep_size et la sous-config pipeline. Indiquez que dp_size est déduit de world_size / (tp_size * pp_size * cp_size).

distributed:
  strategy: fsdp2
  tp_size: 8
  pp_size: 4
  cp_size: 1
  ep_size: 1
  pipeline:
    pp_schedule: interleaved1f1b
    pp_microbatch_size: 1

Parallélisme d'experts MoE

Recommandez strategy: fsdp2 avec ep_size > 1. Dites que cela crée une moe_mesh séparée ; incluez la sous-config moe le cas échéant ; indiquez que ep_size doit diviser dp_size * cp_size. Ne recommandez pas megatron_fsdp ou ddp.

distributed:
  strategy: fsdp2
  ep_size: 8
  moe:
    reshard_after_forward: false

Limitations de MegatronFSDP

Dites non au parallélisme de pipeline, au parallélisme d'experts et à sequence_parallel. Recommandez fsdp2 pour PP, EP ou sequence_parallel ; mentionnez que DDP est seulement du parallélisme de données simple.

Sélection de stratégie

Trois stratégies sont disponibles, sélectionnées via la clé YAML distributed.strategy :

Arbre de décision :

• GPU unique : aucune config distribuée nécessaire (FSDP2Manager ignore la parallélisation quand world_size=1).
• Multi-GPU sur un seul nœud : fsdp2 (par défaut). Utilisez ddp seulement si vous avez besoin de la configuration la plus simple possible.
• Multi-nœud : fsdp2 avec dimensionnement TP/PP approprié.
• Modèles MoE avec parallélisme d'experts : fsdp2 avec ep_size > 1 (crée une moe_mesh séparée).
• Grands modèles (70B+) : fsdp2 avec PP + TP.
• Longues séquences (8K+) : ajoutez CP (cp_size > 1).

Quand vous répondez à des questions de sélection de stratégie, énoncez d'abord le distributed.strategy choisi, puis énumérez les champs YAML que l'utilisateur doit définir.

Réponse rapide TP + PP :

• Utilisez strategy: fsdp2 ; n'utilisez pas megatron_fsdp quand le parallélisme de pipeline est requis.
• Réglez tp_size pour le parallélisme de tenseur et pp_size pour le parallélisme de pipeline.
• Ajoutez une sous-config pipeline: avec pp_schedule et pp_microbatch_size.
• Laissez dp_size non défini ou none ; il est déduit comme world_size / (tp_size * pp_size * cp_size).
• Gardez TP dans un domaine intra-nœud rapide quand possible et utilisez PP à travers la profondeur du modèle pour les modèles 70B+.

Réponse rapide MoE expert-parallel :

• Commencez avec strategy: fsdp2 et ep_size > 1.
• Incluez une sous-config moe: uniquement quand ep_size > 1 ; elle mappe à MoEParallelizerConfig.
• Attendez une moe_mesh séparée pour le parallélisme d'experts en plus de la device_mesh principale.
• Ne recommandez pas megatron_fsdp ou ddp pour le parallélisme d'experts ; megatron_fsdp n'a pas le support EP.
• Avant de terminer une réponse MoE EP, énoncez explicitement que ep_size doit diviser dp_size * cp_size et que megatron_fsdp ne supporte pas EP, PP ou sequence_parallel.

Structure de config YAML

La section distributed dans la recipe YAML mappe directement à parse_distributed_section() dans recipes/_dist_setup.py :

distributed:
  strategy: fsdp2           # fsdp2 | megatron_fsdp | ddp
  dp_size: none             # auto-calculé à partir de world_size / (tp * pp * cp)
  dp_replicate_size: none   # FSDP2 seulement, pour HSDP
  tp_size: 1
  pp_size: 1
  cp_size: 1
  ep_size: 1

  # Drapeaux spécifiques à la stratégie (transmis à la dataclass de stratégie) :
  sequence_parallel: false
  activation_checkpointing: false
  defer_fsdp_grad_sync: true   # FSDP2 seulement

  # Sous-configs (optionnel) :
  pipeline:
    pp_schedule: 1f1b
    pp_microbatch_size: 1
    # ... voir les champs de PipelineConfig

  moe:
    reshard_after_forward: false
    # ... voir les champs de MoEParallelizerConfig

Le dp_size est toujours déduit :

dp_size = world_size / (tp_size * pp_size * cp_size)

Flux d'infrastructure

Section distributed YAML
    -> parse_distributed_section()          [recipes/_dist_setup.py]
    -> setup_distributed()                  [recipes/_dist_setup.py]
        -> create_device_mesh()             [components/distributed/device_mesh.py]
        -> MeshContext(...)                  [components/distributed/mesh.py]
    -> instantiate_infrastructure()         [_transformers/infrastructure.py]
        -> _instantiate_distributed()       -> FSDP2Manager / MegatronFSDPManager / DDPManager
        -> _instantiate_pipeline()          -> AutoPipeline (si pp_size > 1)
        -> parallelize_fn                   -> MoE parallelizer (si ep_size > 1) ou PP wrapper
    -> apply_model_infrastructure()         [_transformers/infrastructure.py]
        -> _shard_pp() ou _shard_ep_fsdp()  (applique le sharding au modèle)

Configuration FSDP2

FSDP2 basique (parallélisme de données seulement)

distributed:
  strategy: fsdp2
  tp_size: 1
  cp_size: 1

Cela auto-calcule dp_size = world_size et applique fully_shard() par bloc transformer via le sharding basé sur DTensor.

FSDP2 avec parallélisme de tenseur

Gardez TP au sein d'un seul domaine NVLink (typiquement un nœud) :

distributed:
  strategy: fsdp2
  tp_size: 4        # 2, 4, ou 8 -- doit diviser les GPUs par nœud
  sequence_parallel: true

Le plan TP est auto-sélectionné selon le type de modèle. Passez un plan personnalisé via l'API Python si nécessaire :

config = FSDP2Config(sequence_parallel=True, tp_plan=my_custom_plan)

FSDP2 avec parallélisme de pipeline

distributed:
  strategy: fsdp2
  pp_size: 2
  pipeline:
    pp_schedule: interleaved1f1b   # 1f1b, gpipe, interleaved_1f1b, etc.
    pp_microbatch_size: 4
    scale_grads_in_schedule: false

Le modèle doit avoir un attribut _pp_plan (défini sur la classe du modèle HF) pour qu'AutoPipeline sache comment diviser les couches à travers les étapes. Les modèles sans _pp_plan ne sont pas compatibles avec PP.

FSDP2 avec HSDP (Hybrid Sharded Data Parallel)

Sharding complet intra-nœud + réplication inter-nœud via une DeviceMesh 2D :

distributed:
  strategy: fsdp2
  dp_replicate_size: 2   # doit diviser dp_size

Contrainte : dp_replicate_size < dp_size (la pure réplication sans sharding n'est pas supportée par FSDP2).

Checkpointing d'activation

Échange du calcul pour la mémoire en recalculant les activations lors du backward :

distributed:
  activation_checkpointing: true

Ceci est transmis à la config de stratégie pour les modèles non-EP, ou lu depuis MeshContext.activation_checkpointing pour les modèles EP.

Deferral de synchronisation de gradient

FSDP2 reporte par défaut la synchronisation de gradient au dernier micro-batch pour le chevauchement de communication :

distributed:
  defer_fsdp_grad_sync: true   # par défaut

Précision mixte

FSDP2Config utilise par défaut bfloat16 pour les trois paramètres de précision via MixedPrecisionPolicy(param_dtype=bf16, reduce_dtype=bf16, output_dtype=bf16, cast_forward_inputs=True). Remplacez via l'API Python :

from torch.distributed.fsdp import MixedPrecisionPolicy
config = FSDP2Config(
    mp_policy=MixedPrecisionPolicy(param_dtype=torch.float16, reduce_dtype=torch.float32),
)

Parallélisme de pipeline

Prérequis

1. La classe de modèle doit définir _pp_plan (un dict mappant les FQN de modules aux étapes).
2. pp_size > 1 dans la section distribuée.
3. Une sous-config pipeline avec schedule et taille de microbatch.

Schedules supportés

Définis dans PipelineConfig.pp_schedule :

• 1f1b (one-forward-one-backward, par défaut)
• gpipe
• interleaved_1f1b / interleaved1f1b
• looped_bfs
• dfs
• v_schedule
• zero_bubble

Exemple (modèle 8B sur 8 GPUs, PP=2 + DP=4)

distributed:
  strategy: fsdp2
  pp_size: 2

  pipeline:
    pp_schedule: interleaved1f1b
    pp_microbatch_size: 4
    scale_grads_in_schedule: false

checkpoint:
  model_save_format: safetensors
  save_consolidated: true

Comment cela fonctionne

AutoPipeline.build() appelle pipeline_model() qui divise le modèle en étapes en utilisant le _pp_plan du modèle, crée des objets PipelineStage et construit le schedule. Pendant l'entraînement, schedule.step() pilote forward et backward à travers le pipeline.

Parallélisme de contexte

Utilisez CP pour les longues séquences (8K+). CP fragmente Q/K/V sur la dimension de séquence sous forme de DTensors.

Config

distributed:
  strategy: fsdp2
  cp_size: 2   # ou 4, 8

Prérequis

• SDPA (Flash Attention ou Efficient Attention backend) ou attention Transformer Engine. SDPBackend.MATH n'est pas compatible avec DTensor.
• Les masques d'attention sont automatiquement supprimés ; is_causal=True est défini via des forward pre-hooks enregistrés par attach_context_parallel_hooks().

Comment cela fonctionne

1. Après le sharding du modèle, apply_model_infrastructure() appelle attach_context_parallel_hooks() sur chaque partie du modèle (pour les modèles non-TE).
2. À chaque étape d'entraînement, make_cp_batch_and_ctx() crée un gestionnaire de contexte CP qui fragmente le batch le long de la dimension de séquence et configure context_parallel() depuis torch.distributed.tensor.experimental.
3. Pour les modèles d'attention TE, make_cp_batch_for_te() utilise le format THD et la fonction thd_get_partitioned_indices de TE pour le fragmentation.

CP avec sequence packing

CP fonctionne avec les séquences empaquetées. Le packed_sequence_size doit être divisible par cp_size. Quand vous utilisez TE, les chunks sont fragmentés par chunk via _shard_thd_chunk_for_te().

Sequence Packing

Emballage de plusieurs séquences dans un seul échantillon d'entraînement pour l'efficacité.

Config

packed_sequence:
  packed_sequence_size: 4096   # 0 = désactivé

step_scheduler:
  local_batch_size: 1          # doit être 1 pour les séquences empaquetées

Quand packed_sequence_size > 0, le collator du dataset empaquette les séquences jusqu'à cette longueur. local_batch_size doit être 1 parce que chaque « sample » est déjà un batch emballé.

Entraînement distribué MoE

Parallélisme d'experts

Réglez ep_size > 1 pour distribuer les experts à travers les GPUs. Cela crée une moe_mesh séparée aux côtés de la device_mesh principale :

distributed:
  strategy: fsdp2
  ep_size: 8
  activation_checkpointing: true

La forme de moe_mesh est (pp_size, ep_shard_size, ep_size) avec des noms de dimension ("pp", "ep_shard", "ep").

Contrainte : dp_cp_size (= dp_size * cp_size) doit être divisible par ep_size.

Sous-config MoE

distributed:
  strategy: fsdp2
  ep_size: 8
  activation_checkpointing: true

  moe:
    reshard_after_forward: false
    ignore_router_for_ac: false
    wrap_outer_model: true

La sous-section moe mappe à MoEParallelizerConfig et n'est instantiée que quand ep_size > 1.

Exemple MoE complet (Qwen3-30B-A3B sur 8 GPUs)

distributed:
  strategy: fsdp2
  tp_size: 1
  cp_size: 1
  pp_size: 1
  ep_size: 8
  sequence_parallel: false
  activation_checkpointing: true

Limitations de MegatronFSDP

Malgré son nom, megatron_fsdp ne supporte pas le parallélisme d'experts (ep_size > 1), le parallélisme de pipeline (pp_size > 1), ou sequence_parallel. Utilisez fsdp2 pour ces fonctionnalités.

Directives de dimensionnement du parallélisme

Modèles denses

Modèles MoE

Les modèles MoE ont besoin de moins de TP que les modèles denses de nombre de paramètres total similaire parce que seulement une fraction des paramètres sont actifs par token. EP est la dimension d'échelle principale :

Règles de topologie matérielle

• TP doit rester au sein d'un seul domaine NVLink (un nœud, typiquement 8 GPUs).
• Utilisez PP ou DP pour l'échelle multi-nœud.
• TP à travers InfiniBand dégrade sévèrement le throughput.

Code Anchors

• components/distributed/config.py : FSDP2Config, MegatronFSDPConfig, DDPConfig.
• components/distributed/mesh.py : MeshContext, strategy map et tailles de maille.
• components/distributed/device_mesh.py : création de device mesh et moe_mesh.
• components/distributed/pipelining/config.py : champs de PipelineConfig.
• components/moe/config.py : MoEParallelizerConfig et MoEConfig.
• recipes/_dist_setup.py : parsing YAML et configuration distribuée.

Pièges

1. TP à travers les nœuds détruit le throughput. Gardez toujours TP au sein d'un seul domaine NVLink. Utilisez PP ou DP pour l'échelle multi-nœud.

2. PP require _pp_plan sur la classe de modèle. Pas tous les modèles HF l'ont. Vérifiez validate_hf_model_for_pipeline_support() avant d'activer PP.

3. Les bulles PP réduisent l'utilisation des GPUs. Utilisez des schedules entrelacés (interleaved_1f1b) et des microbatches plus petits pour réduire le temps de bulle.

4. FSDP2 requires DTensor-aware state dict saving. Utilisez safetensors avec save_consolidated: true pour la compatibilité des checkpoints.

5. CP requires compatible attention. SDPA (Flash Attention ou Efficient Attention) ou attention TE seulement. SDPBackend.MATH n'est pas compatible avec DTensor.

6. *La taille de MoE EP doit diviser uniformément `dp_size cp_size.** La création de device mesh affirmedp_cp_size % ep_size == 0`.

7. MegatronFSDP est plus limité que FSDP2. Il ne supporte pas PP (pp_size > 1), EP (ep_size > 1), ou sequence_parallel. La validation de MeshContext lève une erreur sur ces combinaisons.

8. DDP ne supporte rien au-delà du parallélisme de données. Pas de TP, PP, CP, EP ou HSDP. La validation lève sur n'importe lequel de ceux-ci.

9. Activation checkpointing augmente le calcul. Il économise la mémoire en recalculant les activations lors du backward, mais ajoute ~30% de surcharge de calcul.

10. Mixed precision policy doit correspondre aux attentes du modèle. La policy bfloat16 par défaut fonctionne pour la plupart des modèles. Les modèles FP16 peuvent nécessiter une MixedPrecisionPolicy personnalisée.

11. packed_sequence_size doit être divisible par cp_size quand on utilise CP avec des séquences empaquetées.

12. dp_replicate_size est FSDP2 seulement. Le passer avec megatron_fsdp ou ddp lève une ValueError.

Vérification

Lancez la plus petite recipe qui exerce la stratégie demandée. Le succès signifie code de sortie 0, loss fini, aucun timeout NCCL, et sortie log correspondant aux tailles TP/PP/CP/EP attendues.