Graphes CUDA

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

Qu'est-ce que c'est

Les graphes CUDA capturent les opérations GPU une fois et les rejouent avec une surcharge minale hôte-pilote. Bridge supporte deux implémentations :

Décision rapide

Commencez par des graphes TE-scoped pour la plupart des charges de travail d'entraînement :

• modèles denses : attn, puis optionnellement mlp
• MoE sans droppage : attn moe_router moe_preprocess
• VLMs : la même portée dropless-MoE, mais seulement après que le chemin real-data soit stable

Utilisez local + full_iteration uniquement si vous voulez spécifiquement une capture full-iteration et pouvez satisfaire les contraintes plus strictes.

Pour les charges de travail intensives en recompute :

• les graphes TE-scoped s'associent naturellement au selective recompute
• le full recompute vous pousse généralement vers des graphes local full-iteration ou à l'écart 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 TE-scoped (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 TE-scoped (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 harnais de performance

python scripts/performance/run_performance_workload.py \
  --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"]

Contraintes requises

• use_te_rng_tracker = True (appliqué dans gpt_provider.py)
• portée full_iteration uniquement avec cuda_graph_impl = "local"
• 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 tenseurs 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, réglez NCCL_GRAPH_REGISTER=0 (MCore l'applique pour l'impl local sur arch < sm_100 ; l'impl TE affirme inconditionnellement)
• CPU offloading est incompatible avec les graphes CUDA
• la portée moe_preprocess requiert que la portée moe_router soit également définie

Ordre pratique de déploiement

1. Stabilisez d'abord l'exécution eager.
2. Corrigez la longueur de séquence et la taille du micro-batch.
3. Activez la portée de graphe la plus étroite utile.
4. Confirmez que la relecture est active et que la mémoire reste acceptable.
5. Seulement ensuite élargissez la portée ou combinez avec des fonctionnalités de chevauchement.

Ancres de code

Configuration et validation Bridge

        # CUDA graph scope validation: check_for_nan_in_loss must be disabled with full_iteration graph
        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 must be disabled when using full_iteration CUDA graph. "
                "Set rerun_state_machine.check_for_nan_in_loss=False."
            )
        if self.model.cuda_graph_impl == "none":
            self.model.cuda_graph_scope = []

Exigence TE RNG tracker

        if self.cuda_graph_impl != "none":
            assert getattr(self, "use_te_rng_tracker", False), (
                "Transformer engine's RNG tracker is required for cudagraphs, it can be "
                "enabled with use_te_rng_tracker=True'."

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

    # Capture 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 préchauffage

        # Capture CUDA Graphs after 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 delayed wgrad + graphe CUDA

            cuda_graph_scope = getattr(model_cfg, "cuda_graph_scope", []) or []
            # ... scope parsing ...
            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, ...

Aide helper de remplacement du harnais perf

def _set_cuda_graph_overrides(
    recipe, cuda_graph_impl=None, cuda_graph_scope=None
):
    # Sets impl, scope, and auto-enables te_rng_tracker

Nettoyage de graphe

def _delete_cuda_graphs(cuda_graph_helper):
    # Deletes FullCudaGraphWrapper and TE graph objects to free NCCL buffers

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
• énumération CudaGraphScope: megatron/core/transformer/enums.py

Ancres de recettes positives

• 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. TE RNG tracker est obligatoire : Définir cuda_graph_impl sans use_te_rng_tracker=True et rng.te_rng_tracker=True affirme dans le fournisseur.

2. full_iteration requiert les vérifications NaN désactivées : La totalité de fwd+bwd est capturée, donc la vérification loss-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 token-dropless ne peut graphiquer que moe_router et moe_preprocess, pas la totalité du dispatching expert.

4. Surcharge mémoire : les graphes CUDA épinglent tous les buffers intermédiaires pour la durée de vie du graphe (aucune réutilisation de mémoire). Les graphes TE-scoped ajoutent quelques Go ; les graphes full-iteration peuvent augmenter la mémoire de pointe de 1,5–2×. PP > 1 compose la surcharge puisque chaque étage tient son propre graphe.

5. Interaction delayed wgrad : 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 aucun 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 packed padded si le packing est nécessaire.

7. Le nettoyage des graphes est requis : les objets graphe CUDA tiennent les références de buffer NCCL. Bridge gère cela 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), réglez NCCL_GRAPH_REGISTER=0 lors de l'utilisation de PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. Appliqué dans MCore CudaGraphManager (cuda_graphs.py:1428) et TECudaGraphHelper (cuda_graphs.py:1697). L'impl TE affirme inconditionnellement quelle que soit l'arch.

9. CPU offloading incompatible : les graphes CUDA ne peuvent pas être utilisés avec CPU offloading. Appliqué dans MCore transformer_config.py:1907.

10. Recompute MoE + portée moe_router : le recompute MoE n'est pas supporté avec la portée graphe CUDA moe_router lors de l'utilisation de cuda_graph_impl = "transformer_engine". Appliqué dans MCore transformer_config.py:1977.

11. Recompute au niveau de la couche requiert la portée full_iteration : utiliser recompute_granularity="full" avec recompute_num_layers (recompute N couches transformer entières) est incompatible avec les graphes TE-scoped. MCore appelle cela la granularité "full" même si vous sélectionnez combien de couches — le nom fait référence au recompute de la couche complète, pas du modèle complet. Toute portée TE-scoped (attn, mlp, moe_router, etc.) affirmera : AssertionError: full recompute is only supported with full iteration CUDA graph. Cela touche généralement les configs FP8 qui définissent par défaut les graphes TE-scoped (p. ex. LLAMA3_70B_SFT_CONFIG_H100_FP8_CS_V1 utilise cuda_graph_impl="transformer_engine", cuda_graph_scope="mlp"). Correction : utilisez le recompute de sous-module (recompute_granularity="selective" + recompute_modules), désactivez les graphes CUDA, ou passez à local + full_iteration. Appliqué dans MCore transformer_config.py:2001-2005. Voir aussi @skills/perf-activation-recompute/SKILL.md.

12. Les chiffres de benchmark sont spécifiques à la charge de travail : les gains graphe sont généralement réels quand la surcharge hôte est visible, mais le gain exact dépend de la forme du batch, de la profondeur PP, du recompute, et du fait que la baseline eager était déjà optimisée.

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 de fumée 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 graphe CUDA.
• Aucune erreur NCCL ou accès mémoire illégal dans les logs.