Déchargement CPU

Références

• Documentation stable : @docs/training/cpu-offloading.md
• Métadonnées structurées : @skills/perf-cpu-offloading/card.yaml

Qu'est-ce que c'est

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

Aide à la décision rapide

Activation

Déchargement CPU de l'optimiseur (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

Remplacements CLI :

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

Déchargement CPU d'activation (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'activation

Compatibilité et contraintes

Déchargement d'activation

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

Déchargement d'optimiseur

• Requiert use_distributed_optimizer = True (défaut dans la plupart des recettes)
• Pas de restrictions PP, recompute ou graphe CUDA
• optimizer_offload_fraction doit être dans [0.0, 1.0]

Pratique : grands modèles MoE

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

Config minimale fonctionnelle

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

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

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

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

Déchargement d'activation (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 de 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

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

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

• La validation de config réussit pour le mode de déchargement sélectionné
• L'entraînement se termine sans erreur OOM ou NCCL
• La loss correspond à la baseline non déchargée (delta max < 0.001)
• L'utilisation mémoire baisse proportionnellement à la fraction de déchargement

Ancres de code

Contraintes de déchargement d'activation 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é des graphes CUDA MCore

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

Exclusion mutuelle du déchargement fin 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 de graphe CUDA Bridge

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

Déchargement d'activation 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 de déchargement d'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 de défaillance

Limitations connues

• Le déchargement d'activation requiert PP=1, le rendant impractique pour les grands modèles (30B+ MoE) qui ont besoin du parallélisme de pipeline.
• La pénalité de débit du déchargement d'optimiseur s'ajuste linéairement (~1,9x à 25%, ~4,2x à 100% pour Qwen3-30B-A3B).
• Le chevauchement D2H/H2D fournit seulement ~7% d'accélération car le calcul CPU Adam est le goulot d'étranglement dominant.
• fine_grained_activation_offloading est une approche au niveau module séparée qui fonctionne avec PP > 1 mais ne peut pas être combinée avec le déchargement au niveau couche cpu_offloading.