Déchargement CPU

Références

• Docs stables : @docs/training/cpu-offloading.md
• Métadonnées structurées : @skills/nemo-mbridge-perf-cpu-offloading/card.yaml

Description

Deux mécanismes indépendants pour déplacer les données de la mémoire GPU vers la mémoire CPU :

Prise de décision rapide

Activation

Déchargement optimiseur CPU (recommandé pour les grands modèles)

cfg.optimizer.optimizer_cpu_offload = True
cfg.optimizer.optimizer_offload_fraction = 1.0
cfg.optimizer.overlap_cpu_optimizer_d2h_h2d = True

Surcharges CLI :

optimizer.optimizer_cpu_offload=True \
optimizer.optimizer_offload_fraction=0.5 \
optimizer.overlap_cpu_optimizer_d2h_h2d=True

Déchargement activations CPU (petits/moyens modèles uniquement)

cfg.model.cpu_offloading = True
cfg.model.cpu_offloading_num_layers = 16
cfg.model.cpu_offloading_activations = True
cfg.model.cpu_offloading_weights = False

cfg.model.pipeline_model_parallel_size = 1
cfg.model.recompute_granularity = None
cfg.model.cuda_graph_impl = "none"

Référence des paramètres de config

Déchargement d'optimiseur

Déchargement d'activations

Compatibilité et contraintes

Déchargement d'activations

• pipeline_model_parallel_size doit être 1
• recompute_granularity doit être None
• Impossible de combiner avec fine_grained_activation_offloading
• Impossible de combiner avec CUDA graphs
• cpu_offloading_num_layers doit être dans [0, num_layers-1)

Déchargement d'optimiseur

• Nécessite use_distributed_optimizer = True (défaut dans la plupart des recipes)
• Aucune restriction PP, recompute, ou CUDA graph
• optimizer_offload_fraction doit être dans [0,0, 1,0]

Pratique : grands modèles MoE

Le déchargement d'activations est bloqué pour Qwen3-30B-A3B et modèles MoE similaires. La contrainte PP=1 signifie que chaque GPU retient les 48 couches ; poids du modèle + états optimiseur seuls (~70 GB) dépassent la capacité H100 de 80 GB.

Config minimale fonctionnelle

Déchargement optimiseur (50 %, équilibré)

cfg.optimizer.optimizer_cpu_offload = True
cfg.optimizer.optimizer_offload_fraction = 0.5

Déchargement optimiseur (100 % + chevauchement, max d'économies)

cfg.optimizer.optimizer_cpu_offload = True
cfg.optimizer.optimizer_offload_fraction = 1.0
cfg.optimizer.overlap_cpu_optimizer_d2h_h2d = True

Déchargement activations (petit modèle, PP=1)

cfg.model.cpu_offloading = True
cfg.model.cpu_offloading_num_layers = 16
cfg.model.cpu_offloading_activations = True
cfg.model.cpu_offloading_weights = False
cfg.model.pipeline_model_parallel_size = 1
cfg.model.recompute_granularity = None

Déchargement poids uniquement (petit modèle, PP=1)

cfg.model.cpu_offloading = True
cfg.model.cpu_offloading_num_layers = 8
cfg.model.cpu_offloading_activations = False
cfg.model.cpu_offloading_weights = True
cfg.model.pipeline_model_parallel_size = 1
cfg.model.recompute_granularity = None

Activations et poids (petit modèle, PP=1)

cfg.model.cpu_offloading = True
cfg.model.cpu_offloading_num_layers = 8
cfg.model.cpu_offloading_activations = True
cfg.model.cpu_offloading_weights = True
cfg.model.pipeline_model_parallel_size = 1
cfg.model.recompute_granularity = None

Déchargement poids et déchargement activations partagent les mêmes contraintes (PP=1, pas de recompute, pas CUDA graphs). Déchargement poids n'a pas été testé dans expériences Qwen3-30B-A3B — résultats mesurés couvrent déchargement optimiseur uniquement.

Commande minimale exécutable

uv run python scripts/training/run_recipe.py \
  --recipe qwen3_30b_a3b_pretrain_config \
  optimizer.optimizer_cpu_offload=True \
  optimizer.optimizer_offload_fraction=0.5 \
  train.train_iters=20 \
  train.global_batch_size=8 \
  train.micro_batch_size=1

Vérification

Tests unitaires

uv run python -m pytest \
  tests/unit_tests/models/test_gpt_full_te_layer_autocast_spec.py -k "cpu_offload" \
  tests/unit_tests/peft/test_utils.py -k "cpu_offload" -q

Critères de succès

• Validation config réussit pour mode déchargement sélectionné
• Entraînement se termine sans erreurs OOM ou NCCL
• Loss correspond baseline non-déchargé (delta max < 0,001)
• Utilisation mémoire diminue proportionnellement à fraction déchargement

Ancres de code

Contraintes déchargement activations MCore

        if self.cpu_offloading and (
            self.cpu_offloading_num_layers < 0 or self.cpu_offloading_num_layers >= self.num_layers
        ):
            raise ValueError(...)

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

        if self.cpu_offloading and self.recompute_granularity is not None:
            raise ValueError(
                "CPU offloading does not work when activation recomputation is enabled"
            )

Incompatibilité CUDA graphs MCore

            if self.cpu_offloading:
                raise ValueError("CUDA graphs not supported with CPU offloading.")

Exclusion mutuelle déchargement fin-grained MCore

        if self.fine_grained_activation_offloading:
            assert (
                not self.cpu_offloading
            ), "fine_grained_activation_offloading cannot be enabled with cpu_offloading."

Instanciation HybridDeviceOptimizer MCore

        if config.optimizer_cpu_offload:
            # ... setup cpu/gpu optimizer classes ...
            optimizer = HybridDeviceOptimizer(
                param_groups,
                offload_fraction=config.optimizer_offload_fraction,
                cpu_optimizer_cls=cpu_optimizer_cls,
                gpu_optimizer_cls=gpu_optimizer_cls,
                overlap_cpu_optimizer_d2h_h2d=config.overlap_cpu_optimizer_d2h_h2d,
                pin_cpu_grads=config.pin_cpu_grads,
                pin_cpu_params=config.pin_cpu_params,
            )

Garde CUDA graph Bridge

        assert not config.cpu_offloading and config.recompute_granularity is None, "Cudagraphs not supported"

Déchargement activations Bridge en PEFT

        if self.config.cpu_offloading and self.config.cpu_offloading_activations:
            x.activation_offloading = True
        x, _ = self.linear_in(x)
        x = self.activation(x)
        if self.config.cpu_offloading and self.config.cpu_offloading_activations:
            x.activation_offloading = True
        x, _ = self.linear_out(x)

Champs model_parallel_config MCore

    cpu_offloading: bool = False
    cpu_offloading_num_layers: int = 0
    cpu_offloading_activations: bool = True
    cpu_offloading_weights: bool = False
    cpu_offloading_double_buffering: bool = False
    cpu_offloading_retain_pinned_cpu_buffers: bool = False

Config déchargement optimiseur MCore

    optimizer_cpu_offload: bool = False
    optimizer_offload_fraction: float = 0.0
    use_torch_optimizer_for_cpu_offload: bool = False
    overlap_cpu_optimizer_d2h_h2d: bool = False

Diagnostic d'erreurs

Limitations connues

• Déchargement activations nécessite PP=1, rendant impratique pour grands modèles (30B+ MoE) nécessitant parallelisme pipeline.
• Pénalité débit déchargement optimiseur s'échelonne linéairement (~1,9x à 25 %, ~4,2x à 100 % pour Qwen3-30B-A3B).
• Chevauchement D2H/H2D fournit seulement ~7 % d'accélération car calcul Adam CPU est goulot dominant.
• fine_grained_activation_offloading est approche module-level séparate fonctionnant avec PP > 1 mais impossible à combiner avec cpu_offloading level-couche.