Optimisation mémoire

Docs stables : @docs/parallelisms.md Card : @skills/nemo-mbridge-perf-memory-tuning/card.yaml

Ce que c'est

Les défaillances de dépassement de mémoire GPU (OOM) lors de l'entraînement proviennent souvent d'une fragmentation mémoire plutôt qu'd'une capacité insuffisante. L'allocateur CUDA par défaut de PyTorch peut laisser des lacunes inutilisables entre les allocations. Le correctif le plus efficace est :

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Cela indique à PyTorch d'utiliser des segments mémoire extensibles (de taille non fixe), ce qui réduit dramatiquement la fragmentation et élimine souvent les OOM limites sans aucun changement de modèle ou de parallélisme.

Au-delà de la fragmentation, le pic de mémoire réel est déterminé par :

• Mémoire des paramètres + état du optimiseur — contrôlée par le sharding TP, PP, DP (optimiseur distribué, FSDP)
• Mémoire d'activation — contrôlée par le recompute d'activation, la longueur de séquence, la taille de micro-batch
• Mémoire temporaire / workspace — kernels CUDA, buffers NCCL, graphes CUDA

Pour la planification de la configuration, utilisez l'estimateur théorique de Bridge avant de lancer des travaux volumineux :

from megatron.bridge.training.utils.theoretical_memory_utils import estimate_training_memory

estimate = estimate_training_memory(cfg, num_microbatches=num_microbatches)

L'estimateur rapporte le shard GPU le plus chargé et sépare les composants dense/embedding, MoE routé et d'activation. Il n'inclut pas la fragmentation de l'allocateur, le workspace CUDA/NCCL, les buffers de graphe CUDA, le déséquilibre de tokens ou le workspace dispatcher, donc validez les configs finales avec les métriques de mémoire à l'exécution.

Décision rapide

Quand une exécution d'entraînement provoque un OOM ou est proche de la limite mémoire :

1. Définissez PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True en premier. Cela corrige l'OOM induit par la fragmentation sans coût de performance. La plupart des modèles de lancement Slurm l'incluent déjà.
2. Ajoutez le recompute sélectif d'activation (recompute_modules=[core_attn]) s'il n'est pas déjà activé. Voir @skills/nemo-mbridge-perf-activation-recompute/SKILL.md.
3. Évitez d'augmenter TP comme correctif mémoire — doubler TP augmente dramatiquement le volume d'all-reduce NVLink et tue souvent le throughput (-28% sur Llama3 70B).
4. Évitez d'augmenter PP au détriment du DP — diminuer de moitié le DP double les étapes d'accumulation de gradient et nuit au throughput (~6%).
5. Considérez le recompute mlp si toujours OOM. Économise ~3 GB mais coûte ~16% d'utilisation GPU sur les gros modèles denses (Llama3 70B).
6. Le déchargement CPU est bloqué quand PP > 1.

Activation

Segments extensibles (première étape recommandée)

Définissez dans l'environnement du travail avant le lancement :

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Dans les scripts Slurm, c'est généralement placé à côté d'autres variables d'env :

export CUDA_DEVICE_MAX_CONNECTIONS=1
export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Aucune modification de config de modèle nécessaire. Coût de throughput zéro.

Redimensionnement du parallélisme

Si le modèle ne rentre vraiment pas (pas fragmentation), ajustez le parallélisme :

Recompute d'activation

Voir @skills/nemo-mbridge-perf-activation-recompute/SKILL.md pour les détails complets.

Déchargement CPU

cfg.model.cpu_offloading = True

Incompatible avec PP > 1. Utilisable seulement quand pipeline_model_parallel_size = 1.

Une note sur VPP

Le parallélisme de pipeline virtuel (VPP) est principalement une optimisation de throughput qui réduit les frais généraux de bulles pipeline en entrelacant des chunks de modèle plus petits. Son effet sur le pic mémoire est minimal — changer VPP ne change pas significativement la mémoire d'activation, de paramètres ou d'optimiseur totale sur un GPU.

Dans des expériences antérieures, nous avons incorrectement attribué un correctif OOM au tuning VPP (VPP 5→10). Le correctif réel était PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True qui a éliminé la fragmentation mémoire. La run VPP=10 a en fait utilisé légèrement plus de pic mémoire (60,2 GB vs 58,8 GB) mais n'a pas provoqué d'OOM parce que les segments extensibles ont prévenu la fragmentation.

VPP doit être tuné pour la réduction de bulles pipeline (voir @docs/parallelisms.md), pas comme correctif mémoire.

Compatibilité et contraintes

• expandable_segments:True est incompatible avec --use-nccl-ub (enregistrement de buffer utilisateur NCCL). Voir docs Megatron-FSDP.
• Lors de l'utilisation de graphes CUDA avec expandable_segments:True, définissez NCCL_GRAPH_REGISTER=0 (requis sur les GPUs pré-Blackwell, appliqué par CudaGraphManager MCore).
• Le déchargement CPU nécessite pipeline_model_parallel_size = 1.
• L'optimiseur distribué nécessite use_distributed_optimizer = True dans la config de l'optimiseur.

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 d'or : 709,93 TFLOP/s/GPU
• Seuil de régression : 5%

Comparaison des stratégies : changements de parallélisme pour réduction mémoire

Points clés :

• expandable_segments:True est le gagnant. Le baseline OOM était causé par la fragmentation mémoire, pas une capacité insuffisante. Définir cette variable d'env a éliminé l'OOM sans coût de throughput et sans changement de parallélisme.
• PP=8 fonctionne pour la mémoire mais perd le DP (2→1), signifiant 32 étapes d'accumulation de gradient par batch, ce qui nuit au throughput d'environ 6%.
• TP=8 est catastrophique (-28%) car doubler TP augmente le volume de communication d'all-reduce proportionnellement à travers NVLink, et DP=1 signifie pas de chevauchement de micro-batch.

Déchargement CPU : bloqué

ValueError: Currently there is no support for Pipeline parallelism with CPU offloading. Cette approche est bloquée pour tout modèle utilisant PP > 1.

Recompute d'activation : alternative coûteuse

Le recompute sélectif d'activation avec mlp a économisé ~3 GB de pic mémoire mais a coûté ~16% d'utilisation GPU sur cette charge de travail. Voir @skills/nemo-mbridge-perf-activation-recompute/SKILL.md pour les résultats complets.

Ancres de code

Incompatibilité PP du déchargement 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"
            )

Config VPP et validation de divisibilité de couches (MCore)

            if pipeline_parallel_size and self.virtual_pipeline_model_parallel_size is not None:
                num_layers_per_middle_pipeline_rank = num_layers // pipeline_parallel_size
                if (
                    not num_layers_per_middle_pipeline_rank
                    % self.virtual_pipeline_model_parallel_size
                    == 0
                ):
                    raise ValueError(
                        f"number of layers on each middle pipeline rank:"
                        f"{num_layers_per_middle_pipeline_rank} must be divisible by virtual"
                        f"pipeline parallel degree {self.virtual_pipeline_model_parallel_size}"
                    )

Docs parallélisme sur l'ordonnancement interleaved pipeline

To minimize the pipeline bubble, the computation on each GPU can be divided into multiple subsets of layers (referred to as model chunks), rather than a single contiguous block. Enable this by setting `virtual_pipeline_model_parallel_size`:

model_config = GPTModelProvider(
    pipeline_model_parallel_size=4,
    virtual_pipeline_model_parallel_size=2,  # 2 model chunks per pipeline stage
    # ... other model parameters
)

Diagnostic des défaillances

Limitations connues

• Le déchargement CPU est bloqué quand PP > 1
• Le redimensionnement du parallélisme (TP/PP) a souvent des coûts de throughput significatifs
• L'estimateur théorique est basé sur des formules et ne remplace pas le profilage à l'exécution ou les rapports mémoire CUDA

Vérification

Vérification rapide que expandable_segments:True est actif :

import os
assert "expandable_segments:True" in os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")

Pour les travaux Slurm, vérifiez que la variable d'env est exportée avant la commande d'entraînement dans le script de lancement.