CUDA Graphs

Docs stables : @docs/training/cuda-graphs.md Fiche : @skills/nemo-mbridge-perf-cuda-graphs/card.yaml

De quoi s'agit-il

Les CUDA graphs capturent les opérations GPU une fois et les rejouent avec un surcoût minimal du host-driver. Bridge supporte deux implémentations :

Décision rapide

Commencez par les graphes scoped TE pour la plupart des charges d'entraînement, puis vérifiez le timing de rejeu par rapport à eager sur le même dispatcher, layout et conteneur :

• modèles denses : attn, puis optionnellement mlp
• MoE sans suppression de tokens : attn moe_router moe_preprocess
• VLMs : la même portée MoE sans suppression, mais seulement après stabilisation du chemin données réelles

Utilisez local + full_iteration uniquement si vous voulez spécifiquement capturer l'itération complète et pouvez satisfaire les contraintes plus strictes.

Pour les charges avec beaucoup de recalcul :

• Les graphes scoped TE s'associent naturellement au recalcul sélectif
• Le recalcul complet vous pousse généralement vers des graphes local full-iteration ou loin des graphes entièrement

Docs connexes :

• @docs/training/cuda-graphs.md
• @docs/training/activation-recomputation.md

Activation

Graphe local full-iteration

cfg.model.cuda_graph_impl = "local"
cfg.model.cuda_graph_scope = ["full_iteration"]
cfg.model.cuda_graph_warmup_steps = 3
cfg.model.use_te_rng_tracker = True
cfg.rng.te_rng_tracker = True
cfg.rerun_state_machine.check_for_nan_in_loss = False
cfg.ddp.check_for_nan_in_grad = False

Graphe scoped TE (modèle dense)

cfg.model.cuda_graph_impl = "transformer_engine"
cfg.model.cuda_graph_scope = ["attn"]           # ou ["attn", "mlp"]
cfg.model.cuda_graph_warmup_steps = 3
cfg.model.use_te_rng_tracker = True
cfg.rng.te_rng_tracker = True

Graphe scoped TE (modèle MoE)

cfg.model.cuda_graph_impl = "transformer_engine"
cfg.model.cuda_graph_scope = ["attn", "moe_router", "moe_preprocess"]
cfg.model.cuda_graph_warmup_steps = 3
cfg.model.use_te_rng_tracker = True
cfg.rng.te_rng_tracker = True

CLI du harness de performance

uv run python scripts/performance/run_script.py \
  -m qwen \
  -mr qwen3_30b_a3b \
  --task pretrain \
  -g h100 \
  -c bf16 \
  -ng 16 \
  --cuda_graph_impl transformer_engine \
  --cuda_graph_scope attn,moe_router,moe_preprocess \
  ...

Les valeurs CLI valides se trouvent dans scripts/performance/argument_parser.py :

• VALID_CUDA_GRAPH_IMPLS : ["none", "local", "transformer_engine"]
• VALID_CUDA_GRAPH_SCOPES : ["full_iteration", "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba"]

Le harness de performance utilise une valeur --cuda_graph_scope séparée par des virgules et active automatiquement model.use_te_rng_tracker ainsi que rng.te_rng_tracker quand --cuda_graph_impl n'est pas none.

Contraintes requises

• use_te_rng_tracker = True (appliquée dans gpt_provider.py)
• Portée full_iteration uniquement avec cuda_graph_impl = "local"
• La portée full_iteration requiert check_for_nan_in_loss = False
• Ne combinez pas la portée moe et la portée moe_router
• Les formes de tenseur doivent être statiques (seq_length fixe, micro_batch_size fixe)
• Le routage MoE token-dropless limite la portée graphable aux modules denses
• Avec PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, définissez NCCL_GRAPH_REGISTER=0 (MCore l'applique pour l'impl local sur arch < sm_100 ; l'impl TE l'affirme inconditionnellement)
• L'offloading CPU est incompatible avec les CUDA graphs
• La portée moe_preprocess requiert que la portée moe_router soit également définie

Ordre pratique de mise en place

1. Stabilisez d'abord l'exécution eager.
2. Fixez la longueur de séquence et la taille du micro-batch.
3. Activez la portée de graphe la plus étroite et utile.
4. Confirmez que le rejeu est actif et que la mémoire est toujours acceptable.
5. Comparez eager par rapport aux itérations de rejeu de graphe après warmup et capture ; n'incluez pas l'étape de capture dans le timing en régime permanent.
6. Ensuite seulement, élargissez la portée ou combinez avec des fonctionnalités de chevauchement.

Ancrages de code

Config Bridge et validation

        # Validation de portée CUDA graph : check_for_nan_in_loss doit être désactivé avec graphe full_iteration
        if self.model.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in self.model.cuda_graph_scope:
            assert not self.rerun_state_machine.check_for_nan_in_loss, (
                "check_for_nan_in_loss doit être désactivé lors de l'utilisation du graphe CUDA full_iteration. "
                "Définissez rerun_state_machine.check_for_nan_in_loss=False."
            )
        if self.model.cuda_graph_impl == "none":
            self.model.cuda_graph_scope = []

Exigence du tracker RNG TE

        if self.cuda_graph_impl != "none":
            assert getattr(self, "use_te_rng_tracker", False), (
                "Le tracker RNG de Transformer engine est requis pour les cudagraphs, il peut être "
                "activé avec use_te_rng_tracker=True'."

Création et capture de graphe dans la boucle d'entraînement

    # Capturez les CUDA Graphs.
    cuda_graph_helper = None
    if model_config.cuda_graph_impl == "transformer_engine":
        cuda_graph_helper = TECudaGraphHelper(...)
    # ...
    if config.model.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in config.model.cuda_graph_scope:
        forward_backward_func = FullCudaGraphWrapper(
            forward_backward_func, cuda_graph_warmup_steps=config.model.cuda_graph_warmup_steps
        )

Capture de graphe TE après warmup

        # Capturez les CUDA Graphs après warmup.
        if (
            model_config.cuda_graph_impl == "transformer_engine"
            and cuda_graph_helper is not None
            and not cuda_graph_helper.graphs_created()
            and global_state.train_state.step - start_iteration == model_config.cuda_graph_warmup_steps
        ):
            if model_config.cuda_graph_warmup_steps > 0 and should_toggle_forward_pre_hook:
                disable_forward_pre_hook(model, param_sync=False)
            cuda_graph_helper.create_cudagraphs()
            if model_config.cuda_graph_warmup_steps > 0 and should_toggle_forward_pre_hook:
                enable_forward_pre_hook(model)
                cuda_graph_helper.cuda_graph_set_manual_hooks()

Initialisation RNG

        _set_random_seed(
            rng_config.seed,
            rng_config.data_parallel_random_init,
            rng_config.te_rng_tracker,
            rng_config.inference_rng_tracker,
            use_cudagraphable_rng=(model_config.cuda_graph_impl != "none"),
            pg_collection=pg_collection,
        )

Interaction wgrad différé + CUDA graph

            cuda_graph_scope = getattr(model_cfg, "cuda_graph_scope", []) or []
            # ... analyse de portée ...
            if wgrad_in_graph_scope:
                assert is_te_min_version("2.12.0"), ...
                assert model_cfg.gradient_accumulation_fusion, ...
                if attn_scope_enabled:
                    assert not model_cfg.add_bias_linear and not model_cfg.add_qkv_bias, ...

Helper de surcharge du harness perf

def _set_cuda_graph_overrides(
    recipe, cuda_graph_impl=None, cuda_graph_scope=None
):
    # Définit impl, scope, et active automatiquement te_rng_tracker

Nettoyage de graphe

def _delete_cuda_graphs(cuda_graph_helper):
    # Supprime FullCudaGraphWrapper et objets graphe TE pour libérer les buffers NCCL

Classes MCore (dans 3rdparty/Megatron-LM)

• CudaGraphManager : megatron/core/transformer/cuda_graphs.py
• TECudaGraphHelper : megatron/core/transformer/cuda_graphs.py
• FullCudaGraphWrapper : megatron/core/full_cuda_graph.py
• Enum CudaGraphScope : megatron/core/transformer/enums.py

Ancrages de recipe positifs

• scripts/performance/configs/deepseek/deepseek_workload_base_configs.py
• scripts/performance/configs/qwen/qwen3_workload_base_configs.py
• scripts/performance/configs/gpt_oss/gpt_oss_workload_base_configs.py

Tests

Pièges

1. Le tracker RNG TE est obligatoire : Définir cuda_graph_impl sans use_te_rng_tracker=True et rng.te_rng_tracker=True provoquera une assertion dans le provider.

2. full_iteration requiert les vérifications NaN désactivées : L'intégralité de fwd+bwd est capturée, donc la vérification de perte-NaN ne peut pas inspecter les valeurs intermédiaires.

3. Restrictions de portée MoE : Les portées moe et moe_router s'excluent mutuellement. MoE sans suppression de tokens ne peut grapher que moe_router et moe_preprocess, pas la distribution complète d'experts.

4. Surcoût de mémoire : Les CUDA graphs épinglent tous les buffers intermédiaires pendant toute la durée de vie du graphe (aucune réutilisation de mémoire). Les graphes scoped TE ajoutent quelques GB ; les graphes full-iteration peuvent augmenter la mémoire de pointe de 1,5–2×. PP > 1 aggrave le surcoût car chaque étage détient son propre graphe.

5. Interaction wgrad différé : Quand delay_wgrad_compute=True et que l'attention ou le routeur MoE est dans cuda_graph_scope, des contraintes supplémentaires s'appliquent : TE >= 2.12.0, gradient_accumulation_fusion=True, et pas de biais d'attention.

6. Les séquences de longueur variable cassent les graphes : Les longueurs de séquence doivent être constantes entre les étapes. Utilisez des séquences empaquetées avec padding si l'empaquetage est nécessaire.

7. Le nettoyage de graphe est requis : Les objets CUDA graph tiennent des références à des buffers NCCL. Bridge le gère dans _delete_cuda_graphs() à la fin de l'entraînement, mais les sorties anticipées doivent l'appeler explicitement.

8. Architectures GPU plus anciennes : Sur les GPUs avec compute capability < 10.0 (pré-Blackwell), définissez NCCL_GRAPH_REGISTER=0 lors de l'utilisation de PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. Appliquée dans MCore CudaGraphManager (cuda_graphs.py:1428) et TECudaGraphHelper (cuda_graphs.py:1697). L'impl TE l'affirme inconditionnellement indépendamment de l'arch.

9. L'offloading CPU incompatible : Les CUDA graphs ne peuvent pas être utilisés avec l'offloading CPU. Appliquée dans MCore transformer_config.py:1907.

10. Recalcul MoE + portée moe_router : Le recalcul MoE n'est pas supporté avec la portée CUDA graph moe_router lors de l'utilisation de cuda_graph_impl = "transformer_engine". Appliquée dans MCore transformer_config.py:1977.

11. Recalcul au niveau de la couche requiert la portée full_iteration : Utiliser recompute_granularity="full" avec recompute_num_layers (recalculer N couches transformer complètes) est incompatible avec les graphes scoped TE. MCore appelle cela la granularité "full" même si vous sélectionnez combien de couches — le nom fait référence au recalcul de la couche complète, pas du modèle complet. Toute portée scoped TE (attn, mlp, moe_router, etc.) affirme : AssertionError: full recompute is only supported with full iteration CUDA graph. Cela frappe couramment les configs FP8 qui défaut à des graphes scoped TE (p.ex. LLAMA3_70B_SFT_CONFIG_H100_FP8_CS_V1 utilise cuda_graph_impl="transformer_engine", cuda_graph_scope="mlp"). Correctif : utilisez le recalcul de sous-module (recompute_granularity="selective" + recompute_modules), désactivez les CUDA graphs, ou basculez vers local + full_iteration. Appliquée dans MCore transformer_config.py:2001-2005. Voir aussi @skills/nemo-mbridge-perf-activation-recompute/SKILL.md.

12. Les chiffres de benchmark sont spécifiques à la charge : Les gains de graphe sont généralement réels quand le surcoût du host est visible, mais le gain exact dépend de la forme du batch, de la profondeur PP, du recalcul, du backend dispatcher, et si la baseline eager était déjà optimisée.

13. Une capture réussie n'est pas une garantie d'accélération : Le 2026-05-18, Qwen3 30B A3B H100 BF16 pretrain avec le dispatcher all-to-all a capturé avec succès les graphes scoped TE attn,moe_router,moe_preprocess (48 couches graphables, environ 6,9 s de temps de capture sur le rang 0), mais les itérations de rejeu 5-8 ont moyenné 42,00 s contre 41,36 s pour eager. Traitez les graphes scoped comme un candidat de mise en place et validez sur la pile cible.

Vérification

Tests unitaires

uv run python -m pytest \
  tests/unit_tests/training/test_config.py -k "cuda_graph" \
  tests/unit_tests/training/test_comm_overlap.py -k "cuda_graph" \
  tests/unit_tests/models/test_gpt_full_te_layer_autocast_spec.py -k "cuda_graph" -q

Test smoke fonctionnel (requiert GPU)

uv run python -m pytest \
  tests/functional_tests/recipes/test_llama_recipes_pretrain_cuda_graphs.py -q

Critères de succès

• Les tests unitaires passent, couvrant la validation de config pour les deux implémentations local et transformer_engine.
• Le test fonctionnel complète les étapes d'entraînement avec les deux implémentations de CUDA graph.
• Aucune erreur NCCL ou accès illégal à la mémoire dans les logs.