Optimisation de la mémoire

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

De quoi s'agit-il

Les défaillances GPU OOM pendant l'entraînement proviennent souvent d'une fragmentation de la mémoire plutôt qu'd'une capacité insuffisante. L'allocateur CUDA par défaut de PyTorch peut laisser des espaces 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 drastiquement la fragmentation et élimine souvent les OOM limite sans aucun changement de modèle ou de parallélisme.

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

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

Décision rapide

Quand une exécution d'entraînement cause un OOM ou approche la limite mémoire :

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

Activation

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

À définir dans l'environnement du job avant le lancement :

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Dans les scripts Slurm, c'est généralement placé aux côtés d'autres variables d'environnement :

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

Aucun changement de configuration de modèle nécessaire. Coût de débit nul.

Redimensionnement du parallélisme

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

Recalcul d'activation

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

Offloading CPU

cfg.model.cpu_offloading = True

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

Note sur VPP

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

Dans des expériences antérieures, nous avons incorrectement attribué une correction OOM à l'ajustement VPP (VPP 5→10). La correction réelle était PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True qui a éliminé la fragmentation mémoire. L'exécution VPP=10 a en fait utilisé légèrement plus de mémoire de pic (60,2 GB vs 58,8 GB) mais n'a pas causé d'OOM car les segments extensibles ont prévenu la fragmentation.

VPP devrait être ajusté pour la réduction de bulle de pipeline (voir @docs/parallelisms.md), non comme une solution mémoire.

Compatibilité et contraintes

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

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 %

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

Points clés :

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

Offloading 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.

Recalcul d'activation : alternative coûteuse

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

Ancres de code

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

Config VPP et validation de divisibilité de couche (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 horaire de pipeline intercalé

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 de défaillance

Limitations connues

• L'offloading CPU est bloqué quand PP > 1
• Le redimensionnement du parallélisme (TP/PP) a souvent des coûts de débit significatifs
• Aucun profilage mémoire automatique pour recommander la stratégie optimale

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 jobs Slurm, vérifiez que la variable d'environnement est exportée avant la commande d'entraînement dans le script de lancement.