Nemo Mbridge Perf Cpu Offloading

NVIDIAOfficialPublished by NVIDIA

Skill
Tool integration
Software development
Apache-2.0
Medium risk

From the author

Validate and use CPU offloading in Megatron Bridge, including layer-level activation offloading and fractional optimizer state offloading with HybridDeviceOptimizer.

— the description by NVIDIA

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

This skill provides guidance on implementing two distinct CPU offloading mechanisms within Megatron Bridge: activation offloading and optimizer state offloading. These techniques allow developers to move memory-intensive data from GPU to CPU, enabling the training of larger models or fitting models into constrained GPU memory. The skill details configuration parameters, compatibility constraints, and troubleshooting steps for both approaches, including the use of HybridDeviceOptimizer for optimizer states.

What it does

Enables activation offloading to move transformer layer activations to CPU memory.
Supports optimizer state offloading for Adam optimizer states using HybridDeviceOptimizer.
Provides configuration parameters for tuning memory-speed trade-offs.
Includes validation steps to ensure compatibility with pipeline parallelism and CUDA graphs.
Offers troubleshooting tables for common OOM errors and performance bottlenecks.
Provides minimal runnable command examples for training recipes.

When to use it

When GPU memory is insufficient for the model size and you need to offload data to CPU.
When you need to choose between activation offloading (for smaller models) or optimizer offloading (for large models).
When investigating OOM errors or crashes related to memory configuration.
When you need to tune the memory-speed trade-off using fractional optimizer offloading.

Advantages

Reduces GPU memory footprint, allowing for larger model training.
Provides flexible options for different model scales and parallelism strategies.
Includes clear constraints to prevent invalid configuration combinations.
Offers specific guidance on performance impacts and throughput trade-offs.

Usage examples

“How do I enable optimizer CPU offloading for a large model?”: The agent provides the configuration block setting optimizer_cpu_offload to True and configuring the offload fraction and overlap settings.
“Why am I getting an error when using activation offloading with pipeline parallelism?”: The agent explains that activation offloading requires pipeline_model_parallel_size to be 1 and suggests using optimizer offloading instead.

Install

Download the skill as a ZIP:
https://mp-staging.autogineer.com/agent-skills/nvidia/nemo-mbridge-perf-cpu-offloading/download

Unzip it into the cross-client skills folder of your project:

mkdir -p .agents/skills && unzip nemo-mbridge-perf-cpu-offloading.zip -d .agents/skills

Documentation by NVIDIA

CPU Offloading

References

Stable docs: @docs/training/cpu-offloading.md
Structured metadata: @skills/nemo-mbridge-perf-cpu-offloading/card.yaml

What It Is

Two independent mechanisms to move data from GPU to CPU memory:

Mechanism	Config namespace	What gets offloaded	PP restriction
Activation offloading	`model.cpu_offloading*`	Activations (and optionally weights) per transformer layer	PP must be 1
Optimizer offloading	`optimizer.optimizer_cpu_offload`	Adam optimizer states (momentum + variance) via `HybridDeviceOptimizer`	None

Quick Decision

Situation	Recommendation
Large MoE model (30B+), needs PP > 1	Optimizer offloading — activation offloading is blocked by PP=1
Small/medium model, PP=1 fits, activation memory dominates	Activation offloading
Want tunable memory-speed tradeoff	Optimizer offloading with fractional `optimizer_offload_fraction`
Throughput is top priority	Don't enable — offloading always adds overhead
CUDA graphs are needed	Only optimizer offloading — activation offloading is incompatible
Memory pressure is moderate	Optimizer offload at 25–50% fraction for best efficiency

Enablement

Optimizer CPU offloading (recommended for large models)

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

CLI overrides:

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

Activation CPU offloading (small/medium models only)

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"

Config Parameter Reference

Optimizer offloading

Parameter	Default	Description
`optimizer_cpu_offload`	`False`	Master switch
`optimizer_offload_fraction`	`0.0`	Fraction of optimizer states on CPU (0.0–1.0)
`overlap_cpu_optimizer_d2h_h2d`	`False`	Overlap GPU↔CPU transfers with compute
`use_torch_optimizer_for_cpu_offload`	`False`	Use `torch.optim` instead of fused optimizer for CPU portion

Activation offloading

Parameter	Default	Description
`cpu_offloading`	`False`	Master switch
`cpu_offloading_num_layers`	`0`	Number of transformer layers to offload (0 to num_layers-1)
`cpu_offloading_activations`	`True`	Offload activations
`cpu_offloading_weights`	`False`	Offload weights
`cpu_offloading_double_buffering`	`False`	Double-buffer across layers while reloading

Compatibility And Constraints

Activation offloading

pipeline_model_parallel_size must be 1
recompute_granularity must be None
Cannot combine with fine_grained_activation_offloading
Cannot combine with CUDA graphs
cpu_offloading_num_layers must be in [0, num_layers-1)

Optimizer offloading

Requires use_distributed_optimizer = True (default in most recipes)
No PP, recompute, or CUDA graph restrictions
optimizer_offload_fraction must be in [0.0, 1.0]

Practical: large MoE models

Activation offloading is blocked for Qwen3-30B-A3B and similar large MoE models. The PP=1 constraint means each GPU holds all 48 layers; model weights + optimizer states alone (~70 GB) exceed H100 80 GB capacity.

Minimal Runnable Command

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

Verification

Unit tests

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

Success criteria

Config validation passes for the selected offloading mode
Training completes without OOM or NCCL errors
Loss matches the non-offloaded baseline (max delta < 0.001)
Memory usage drops proportionally to offload fraction

Code Anchors

MCore activation offload constraints

        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"
            )

MCore CUDA graph incompatibility

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

MCore fine-grained offloading mutual exclusion

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

MCore HybridDeviceOptimizer instantiation

        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,
            )

Bridge CUDA graph guard

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

Bridge activation offloading in 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)

Failure Diagnosis

Symptom	Likely Cause	How To Confirm	Fix
`Currently there is no support for Pipeline parallelism with CPU offloading`	Activation offload + PP > 1	Check `pipeline_model_parallel_size`	Set PP=1 or use optimizer offloading
`CPU offloading does not work when activation recomputation is enabled`	Activation offload + recompute	Check `recompute_granularity`	Set `recompute_granularity=null`
`fine_grained_activation_offloading cannot be enabled with cpu_offloading`	Both offloading modes enabled	Check both flags	Use one or the other
`CUDA graphs not supported with CPU offloading`	CUDA graphs + activation offload	Check `cuda_graph_impl`	Set `cuda_graph_impl="none"`
OOM with activation offloading	Model too large for PP=1	Check allocated memory vs 80 GB	Use optimizer offloading with PP > 1
Extreme slowdown (>4x)	100% optimizer offload, CPU Adam bottleneck	Compare iter time at different fractions	Reduce fraction or enable `overlap_cpu_optimizer_d2h_h2d`
OOM at partial optimizer offload	Insufficient offload for this config	Check memory at different fractions	Increase fraction or add PP

Known Limitations

Activation offloading requires PP=1, making it impractical for large models (30B+ MoE) that need pipeline parallelism.
Optimizer offloading throughput penalty scales linearly (~1.9x at 25%, ~4.2x at 100% for Qwen3-30B-A3B).
D2H/H2D overlap provides only ~7% speedup because CPU Adam compute is the dominant bottleneck.
fine_grained_activation_offloading is a separate module-level approach that works with PP > 1 but cannot be combined with layer-level cpu_offloading.

Files

ATTRIBUTION.mdAdded by AutoGineer863 B
BENCHMARK.md4.2 KiB
LICENSELicence file added by AutoGineer11.2 KiB
SKILL.md9.1 KiB
card.yaml8.0 KiB
evals/evals.json1.4 KiB
skill-card.md3.8 KiB
skill.oms.sig4.7 KiB

card.yaml

title: cpu_offloading
validated_on: "2026-03-30"
summary: >
  Megatron Bridge supports two CPU offloading mechanisms: layer-level activation
  offloading (cpu_offloading) and fractional optimizer state offloading
  (optimizer_cpu_offload via HybridDeviceOptimizer). Activation offloading requires
  PP=1, no recompute, no CUDA graphs — impractical for large MoE models. Optimizer
  offloading works with PP>1 and supports tunable fraction (0-100%) with linear
  memory-speed tradeoff. Verified on Qwen3-30B-A3B MoE pretrain (TP2 PP2 EP4,
  2 nodes H100): 25-100% offload saves 3.8-15.3 GB (-8% to -32%) at 1.9x-4.2x
  slowdown. All variants numerically safe (loss delta < 0.001). D2H/H2D overlap
  provides ~7% speedup at 100% offload.
validation_status:
  activation_offload_pp_constraint:
    - code_verified  # transformer_config.py:1303-1306
  activation_offload_recompute_constraint:
    - code_verified  # transformer_config.py:1308-1310
  activation_offload_cuda_graph_constraint:
    - code_verified  # transformer_config.py:1943-1944
  activation_offload_fine_grained_constraint:
    - code_verified  # transformer_config.py:1427-1430
  optimizer_offload_hybrid_optimizer:
    - code_verified  # optimizer/__init__.py:480-518
  optimizer_offload_config_fields:
    - code_verified  # optimizer_config.py
  bridge_cuda_graph_guard:
    - code_verified  # gpt_full_te_layer_autocast_spec.py:234
  bridge_peft_offloading:
    - code_verified  # peft/utils.py:621-631
  activation_offload_qwen3_moe:
    - measured  # OOM at 70.3 GB, PP=1 on H100 80 GB. Jobs 10609511, 10609272
  optimizer_offload_25pct:
    - measured  # Qwen3-30B-A3B, TP2 PP2 EP4, 2x H100. Job 10611042
  optimizer_offload_50pct:
    - measured  # Job 10611043
  optimizer_offload_75pct:
    - measured  # Job 10611045
  optimizer_offload_100pct:
    - measured  # Job 10609512
  optimizer_offload_100pct_overlap:
    - measured  # Job 10611046
training_dimensions:
  speed:
    effect: "1.9x-4.2x slower step time (scales linearly with offload fraction)"
    confidence: high
    rationale: >
      CPU Adam compute and D2H/H2D transfers add latency. Measured on
      Qwen3-30B-A3B TP2 PP2 EP4, 2 nodes H100. D2H/H2D overlap reduces
      100% penalty from 4.2x to 3.9x.
  memory:
    effect: "3.8 GB saved per 25% of offload fraction (up to 15.3 GB / 32% at 100%)"
    confidence: high
    rationale: >
      Measured on Qwen3-30B-A3B (47.2 GB baseline). Savings scale linearly.
  scale:
    effect: "enables otherwise-OOM configurations"
    confidence: medium
    rationale: >
      Can free memory for larger batch sizes or additional parallelism.
  convergence:
    effect: "no change (loss delta < 0.001 across all fractions)"
    confidence: high
    rationale: >
      All fractions produce identical loss across 20 iterations on Qwen3-30B-A3B.
  stability:
    effect: "no issues observed"
    confidence: high
    rationale: >
      No errors, hangs, or NCCL issues across 120 total iterations tested.
enable_when:
  - GPU memory is tight and throughput regression is acceptable
  - model requires PP > 1 (use optimizer offloading, not activation offloading)
  - want tunable memory-speed tradeoff via offload fraction
avoid_when:
  - throughput is the primary concern
  - model already fits comfortably in GPU memory
  - CUDA graphs are needed (incompatible with activation offloading)
interactions:
  required:
    - optimizer.use_distributed_optimizer (for optimizer offloading)
  conditional:
    - pipeline_model_parallel_size must be 1 for activation offloading
    - recompute_granularity must be None for activation offloading
  incompatible:
    - cpu_offloading with pipeline_model_parallel_size > 1
    - cpu_offloading with recompute_granularity != None
    - cpu_offloading with fine_grained_activation_offloading
    - cpu_offloading with cuda_graph_impl != none
feature_meaning:
  cpu_offloading: >
    Master switch for layer-level activation CPU offloading.
  optimizer_cpu_offload: >
    Master switch for optimizer state CPU offloading via HybridDeviceOptimizer.
  optimizer_offload_fraction: >
    Fraction of optimizer states on CPU. 0.0 = all GPU, 1.0 = all CPU.
  overlap_cpu_optimizer_d2h_h2d: >
    Overlap GPU-CPU transfers with compute during optimizer step.
config_keys:
  - model.cpu_offloading
  - model.cpu_offloading_num_layers
  - model.cpu_offloading_activations
  - model.cpu_offloading_weights
  - optimizer.optimizer_cpu_offload
  - optimizer.optimizer_offload_fraction
  - optimizer.overlap_cpu_optimizer_d2h_h2d
recommended_path:
  optimizer.optimizer_cpu_offload: true_when_memory_constrained
  optimizer.optimizer_offload_fraction: 0.5_for_balanced_or_1.0_for_max_savings
  optimizer.overlap_cpu_optimizer_d2h_h2d: true_at_high_fractions
  model.cpu_offloading: only_for_small_models_where_pp1_fits
expected_metric_change:
  - metric: memory_allocated
    direction: down
    magnitude: 8-32%
    conditions: optimizer offloading on Qwen3-30B-A3B MoE
    evidence: measured_qwen3_30b_a3b
  - metric: step_time
    direction: up
    magnitude: 1.9x-4.2x
    conditions: optimizer offloading on Qwen3-30B-A3B MoE
    evidence: measured_qwen3_30b_a3b
measured_results:
  - model: Qwen3-30B-A3B
    config: baseline
    iter_ms: 996
    mem_allocated_gb: 47.2
    job: "10119286"
    status: success
  - model: Qwen3-30B-A3B
    config: optimizer_offload_25pct
    iter_ms: 1871
    slowdown: 1.9x
    mem_allocated_gb: 43.4
    mem_saved_gb: 3.8
    loss_match: true
    job: "10611042"
    status: success
  - model: Qwen3-30B-A3B
    config: optimizer_offload_50pct
    iter_ms: 2484
    slowdown: 2.5x
    mem_allocated_gb: 39.6
    mem_saved_gb: 7.6
    loss_match: true
    job: "10611043"
    status: success
  - model: Qwen3-30B-A3B
    config: optimizer_offload_75pct
    iter_ms: 3187
    slowdown: 3.2x
    mem_allocated_gb: 35.6
    mem_saved_gb: 11.6
    loss_match: true
    job: "10611045"
    status: success
  - model: Qwen3-30B-A3B
    config: optimizer_offload_100pct
    iter_ms: 4189
    slowdown: 4.2x
    mem_allocated_gb: 32.0
    mem_saved_gb: 15.3
    loss_match: true
    job: "10609512"
    status: success
  - model: Qwen3-30B-A3B
    config: optimizer_offload_100pct_overlap
    iter_ms: 3877
    slowdown: 3.9x
    mem_allocated_gb: 32.0
    overlap_speedup_pct: 7
    loss_match: true
    job: "10611046"
    status: success
  - model: Qwen3-30B-A3B
    config: activation_offload_24_layers
    mem_allocated_gb: 70.3
    job: "10609511"
    status: oom
failure_modes:
  - name: pp_constraint
    symptom: "Currently there is no support for Pipeline parallelism with CPU offloading"
    fix: set PP = 1 or use optimizer offloading
  - name: recompute_constraint
    symptom: "CPU offloading does not work when activation recomputation is enabled"
    fix: set recompute_granularity = null
  - name: cuda_graph_constraint
    symptom: "CUDA graphs not supported with CPU offloading"
    fix: set cuda_graph_impl = none
  - name: large_model_oom
    symptom: OOM with activation offloading even with many nodes
    fix: use optimizer offloading instead
known_constraints:
  - cpu_offloading requires PP = 1 (transformer_config.py:1303)
  - cpu_offloading requires recompute_granularity = None (transformer_config.py:1308)
  - cpu_offloading incompatible with CUDA graphs (transformer_config.py:1943)
  - optimizer_cpu_offload requires use_distributed_optimizer (optimizer/__init__.py:682)
known_limitations:
  - Activation offloading impractical for large MoE models due to PP=1 constraint
  - Optimizer offloading throughput penalty scales linearly with fraction
  - D2H/H2D overlap provides only ~7% improvement
evidence:
  - docs/training/cpu-offloading.md
  - 3rdparty/Megatron-LM/megatron/core/transformer/transformer_config.py
  - 3rdparty/Megatron-LM/megatron/core/optimizer/optimizer_config.py
  - 3rdparty/Megatron-LM/megatron/core/optimizer/__init__.py
  - src/megatron/bridge/models/gpt_full_te_layer_autocast_spec.py
follow_up_validation:
  - Test optimizer offload + CUDA graphs combined
  - Test fine_grained_activation_offloading as alternative
  - Test with larger GBS to amortize CPU optimizer overhead
  - Test activation offloading on smaller models where PP=1 fits

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Needs Python

Security

What our scanner noticed· Medium risk

Large base64-encoded blobs

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Published by

NVIDIA

Publisher

NVIDIA (Official)

Source

NVIDIA/skills · skills/nemo-mbridge-perf-cpu-offloading · commit fce3cca

Licence

Apache License 2.0 (redistribution allowed)

LICENSE (added by AutoGineer from LICENSE-APACHE)
ATTRIBUTION.md (written by AutoGineer)

Synced

2026-11-08

AutoGineer is not affiliated with or endorsed by NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download nemo-mbridge-perf-cpu-offloading.zip

From the author

AutoGineer summary

What it does

When to use it

Advantages

Usage examples

Install

Documentation by NVIDIA

CPU Offloading

References

What It Is

Quick Decision

Enablement

Optimizer CPU offloading (recommended for large models)

Activation CPU offloading (small/medium models only)

Config Parameter Reference

Optimizer offloading

Activation offloading

Compatibility And Constraints

Activation offloading

Optimizer offloading

Practical: large MoE models

Minimal Runnable Command

Verification

Unit tests

Success criteria

Code Anchors

MCore activation offload constraints

MCore CUDA graph incompatibility

MCore fine-grained offloading mutual exclusion

MCore HybridDeviceOptimizer instantiation

Bridge CUDA graph guard

Bridge activation offloading in PEFT

Failure Diagnosis

Known Limitations

Files

card.yaml

Requirements and compatibility

Security

Attribution and licence

Related skills

Implement Spec

MCP Builder

Schema

Postgresql Table Design

SQL Optimization Patterns

API Scaffolding Fastapi Pro

Download ZIP