Activation Recompute

Stable docs: @docs/training/activation-recomputation.md Card: @skills/perf-activation-recompute/card.yaml

What It Is

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

Note : MCore les nomme « selective » (niveau sous-module) vs « full » (niveau couche). « Full » signifie recalculer des couches complètes, pas le modèle complet — tu choisis toujours le nombre de couches via recompute_num_layers.

Décision rapide

1. Définis PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True d'abord — la plupart des OOMs limites sont causés par la fragmentation mémoire, pas la capacité. Cela le résout à coût zéro. Voir @skills/perf-memory-tuning/SKILL.md.
2. Commence avec recompute_granularity=selective, recompute_modules=[core_attn] (souvent déjà par défaut dans les recettes).
3. Ajoute layernorm aux modules recompute — calcul presque gratuit mais économie mémoire négligeable. Aide seulement dans les cas extrêmement limites.
4. Ajoute mlp en dernier recours — économise ~3 GB mais coûte ~16% d'utilisation GPU sur les grands modèles denses (Llama3 70B).
5. Utilise recompute_granularity=full seulement si le recompute selective ne tient toujours pas.

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

Activation

Selective recompute (défaut pour la plupart des recettes)

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

Selective recompute avec modules supplémentaires

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

Recompute de couche complète

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

Modules recompute disponibles

CLI harnais de performance

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

Compatibilité et contraintes

• recompute_granularity=selective requiert une liste recompute_modules non-vide
• recompute_granularity=full requiert recompute_method et recompute_num_layers
• Le recompute au niveau couche (recompute_granularity="full" + recompute_num_layers) est incompatible avec les CUDA graphs en scope TE. MCore appelle cela la granularité « full » — le nom fait référence au recalcul de couches transformer complètes, pas du modèle complet. Même si tu sélectionnes le nombre de couches à recalculer, MCore le traite différemment du recompute de sous-module. Tout scope en scope TE (attn, mlp, moe_router, etc.) assertera. Cela touche couramment les configs FP8 qui activent les graphes en scope 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 :
  • utilise le recompute de sous-module (recompute_granularity="selective" + recompute_modules) — compatible avec les graphes en scope TE
  • désactive les CUDA graphs (cuda_graph_impl="none") et utilise le recompute au niveau couche
  • bascule vers cuda_graph_impl="local", cuda_graph_scope="full_iteration"
• distribute_saved_activations=True ne peut pas être combiné avec sequence_parallel=True
• Combiner le recompute mlp + core_attn est légèrement pire que mlp seul du fait d'une double surcharge recompute

Résultats mesurés

Llama3 70B SFT sur 32x H100 80GB, FP8 (Mise à l'échelle actuelle) :

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

Éléments clés :

• le recompute layernorm est presque gratuit en calcul mais économise une mémoire négligeable
• le recompute mlp économise ~3 GB pic mais coûte ~16% parce que la FFN Llama3 70B (hidden=28672) est chère à recalculer
• combiner mlp + core_attn est légèrement pire que mlp seul
• pour cette charge de travail, le vrai fix OOM était PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True (fragmentation mémoire, pas capacité). Voir @skills/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 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

Valeurs par défaut recette Llama3

    # Memory saving (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é CPU offloading PP (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 des défaillances

Limitations connues

• L'économie mémoire par module varie significativement selon l'architecture du modèle et la dimension cachée
• Aucune sélection module automatique — les utilisateurs doivent choisir quels modules recalculer
• le recompute layernorm n'en vaut presque jamais la peine comme fix standalone
• le CPU offloading (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 :

• Les tests unitaires passent pour la validation config recompute
• Aucune erreur assertion de validation config