Nemo Mbridge Perf Cpu Offloading — Agent Skill by NVIDIA

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

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.

Implement Spec
Matt Pocock· Community
A workflow for implementing technical specifications and tickets using concurrent subagents.
Software developmentClaude CodeClaude appsOpenAI Codex+14 more
MITGitHub stars of the source repository317.5k
MCP Builder
AnthropicOfficial
A comprehensive guide for building high-quality Model Context Protocol (MCP) servers to integrate external services with LLMs.
Software developmentClaude CodeClaude appsOpenAI Codex+14 more
Apache-2.0Low riskHas scriptsGitHub stars of the source repository188.6k
Schema
Corey Haines· Community
Expert guidance for implementing and optimizing JSON-LD schema markup to improve search engine content understanding.
Software developmentClaude CodeClaude appsOpenAI Codex+14 more
MITGitHub stars of the source repository59k
Postgresql Table Design
Seth Hobson· Community
A comprehensive guide for designing, reviewing, and optimizing PostgreSQL database schemas.
Software developmentClaude CodeClaude appsOpenAI Codex+14 more
MITGitHub stars of the source repository42.3k
SQL Optimization Patterns
Seth Hobson· Community
A comprehensive guide to optimizing SQL database performance through indexing, query plan analysis, and efficient coding patterns.
Software developmentClaude CodeClaude appsOpenAI Codex+14 more
MITGitHub stars of the source repository42.3k
API Scaffolding Fastapi Pro
Seth Hobson· Community
An expert FastAPI developer focused on high-performance, async-first API architecture and production-ready microservices.
Software developmentClaude Code
MITGitHub stars of the source repository42.3k

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

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