Recompute des activations

Docs stables : @docs/training/activation-recomputation.md Fiche : @skills/nemo-mbridge-perf-activation-recompute/card.yaml

Ce que c'est

Le recompute des activations échange du calcul GPU pour de la mémoire en supprimant les activations intermédiaires lors de la passe forward et en les recalculant lors de la backward. Megatron Bridge supporte deux granularités :

Note : MCore nomme cela « selective » (niveau sous-module) vs « full » (niveau couche). « Full » signifie recalculer des couches complètes, pas le modèle entier — vous choisissez toujours combien de couches via recompute_num_layers.

Décision rapide

1. Définissez PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True d'abord — la plupart des OOM limites sont causés par la fragmentation mémoire, pas la capacité. Cela le règle sans coût. Voir @skills/nemo-mbridge-perf-memory-tuning/SKILL.md.
2. Commencez avec recompute_granularity=selective, recompute_modules=[core_attn] (souvent déjà par défaut dans les recettes).
3. Ajoutez layernorm aux modules recompute — calcul presque gratuit mais économies mémoire négligeables. N'aide que dans les cas extrêmement limites.
4. Ajoutez mlp en dernier recours — économise ~3 GB mais coûte ~16% d'utilisation GPU sur les grands modèles denses (Llama3 70B).
5. Utilisez recompute_granularity=full seulement quand le recompute selective ne rentre toujours pas.

Le offloading CPU (cpu_offloading=True) est une alternative qui évite entièrement le coût recompute, mais il est incompatible avec PP > 1.

Activation

Recompute selective (défaut pour la plupart des recettes)

cfg.model.recompute_granularity = "selective"
cfg.model.recompute_modules = ["core_attn"]

Recompute selective avec modules supplémentaires

cfg.model.recompute_granularity = "selective"
cfg.model.recompute_modules = ["core_attn", "layernorm"]  # ou ["mlp"] ou ["mlp", "core_attn"]

Recompute par couche complète

cfg.model.recompute_granularity = "full"
cfg.model.recompute_method = "uniform"
cfg.model.recompute_num_layers = 4

recompute_modules disponibles

CLI du performance harness

python scripts/performance/run_performance_workload.py \
  --recompute_granularity selective \
  --recompute_modules core_attn layernorm \
  ...

Compatibilité et contraintes

• recompute_granularity=selective nécessite une liste recompute_modules non vide
• recompute_granularity=full nécessite recompute_method et recompute_num_layers
• Le recompute au niveau couche (recompute_granularity="full" + recompute_num_layers) est incompatible avec les CUDA graphs scopés TE. MCore appelle cela la granularité « full » — le nom fait référence au recalcul de couches transformer complètes, pas du modèle entier. Même si vous sélectionnez combien de couches recalculer, MCore le traite différemment du recompute sous-module. Tout scope scopé TE (attn, mlp, moe_router, etc.) génèrera une assertion. Cela frappe couramment les configs FP8 qui activent les graphs scopés TE par défaut (ex. LLAMA3_70B_SFT_CONFIG_H100_FP8_CS_V1 définit cuda_graph_impl="transformer_engine", cuda_graph_scope="mlp"). Options :
  • utiliser le recompute sous-module (recompute_granularity="selective" + recompute_modules) — compatible avec les graphs scopés TE
  • désactiver les CUDA graphs (cuda_graph_impl="none") et utiliser le recompute au niveau couche
  • passer à cuda_graph_impl="local", cuda_graph_scope="full_iteration"
• distribute_saved_activations=True ne peut pas être combiné avec sequence_parallel=True
• Combiner mlp + recompute core_attn est légèrement pire que mlp seul en raison du surcoût double recompute

Résultats mesurés

Llama3 70B SFT sur 32x H100 80GB, FP8 (Scaling actuel) :

• Baseline : TP=4, PP=4, VPP=5, DP=2, MBS=1, GBS=32, seq_len=4096
• Utilisation GPU golden : 709,93 TFLOP/s/GPU
• Seuil régression : 5%

Points clés :

• Le recompute layernorm est presque gratuit en calcul mais économise mémoire négligeable
• Le recompute mlp économise ~3 GB pic mais coûte ~16% parce que le FFN Llama3 70B (hidden=28672) est coûteux à recalculer
• Combiner mlp + core_attn est légèrement pire que mlp seul
• Pour cette charge, le vrai correctif OOM était PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True (fragmentation mémoire, pas capacité). Voir @skills/nemo-mbridge-perf-memory-tuning/SKILL.md.

Ancres de code

Enum modules recompute et logique checkpoint selective

# 3rdparty/Megatron-LM/megatron/core/transformer/transformer_block.py
# _checkpointed_forward() applique le recompute selective basé sur recompute_modules

Validation config recompute

# 3rdparty/Megatron-LM/megatron/core/transformer/transformer_config.py
# Valide recompute_granularity, recompute_method, recompute_num_layers

Défauts recette Llama3

    # Économie mémoire (recompute & offloading)
    cfg.model.recompute_granularity = None
    cfg.model.recompute_modules = None
    cfg.model.fine_grained_activation_offloading = False
    cfg.model.offload_modules = None

Recompute full + assertion CUDA graph (MCore)

            if self.recompute_granularity:
                if self.recompute_granularity != "selective":
                    assert self.cuda_graph_scope == [
                        CudaGraphScope.full_iteration
                    ], "full recompute is only supported with full iteration CUDA graph."

Incompatibilité PP offloading CPU (MCore)

        if self.cpu_offloading and self.pipeline_model_parallel_size > 1:
            raise ValueError(
                "Currently there is no support for Pipeline parallelism with CPU offloading"
            )

Diagnostic de défaillance

Limitations connues

• Les économies mémoire par module varient significativement selon l'architecture modèle et la dimension cachée
• Aucune sélection module automatique — les utilisateurs doivent choisir quels modules recalculer
• Le recompute layernorm ne vaut presque jamais le coup en tant que correctif autonome
• Le offloading CPU (l'alternative à coût zéro calcul) est bloqué quand PP > 1

Vérification

uv run python -m pytest \
  tests/unit_tests/training/test_config.py -k "recompute" -q

Critères de succès :

• Tests unitaires passent pour validation config recompute
• Aucune erreur assertion de validation config