Nemo Mbridge Perf Memory Tuning

Memory Tuning

Stable docs: @docs/parallelisms.md Card: @skills/nemo-mbridge-perf-memory-tuning/card.yaml

What It Is

GPU OOM failures during training often stem from memory fragmentation rather than raw capacity. PyTorch's default CUDA allocator can leave unusable gaps between allocations. The single most effective fix is:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

This tells PyTorch to use expandable (non-fixed-size) memory segments, which dramatically reduces fragmentation and often eliminates borderline OOM without any model or parallelism changes.

Beyond fragmentation, actual peak memory is determined by:

• Parameter + optimizer state memory — controlled by TP, PP, DP sharding (distributed optimizer, FSDP)
• Activation memory — controlled by activation recompute, sequence length, micro-batch size, and PEFT-specific retention of gathered inputs
• Temporary / workspace memory — CUDA kernels, NCCL buffers, CUDA graphs

For configuration planning, use the Bridge theoretical estimator before launching large jobs:

from megatron.bridge.training.utils.theoretical_memory_utils import estimate_training_memory

estimate = estimate_training_memory(cfg, num_microbatches=num_microbatches)

The estimator reports the most-loaded GPU shard and separates dense/embedding, routed MoE expert, and activation components. It does not include allocator fragmentation, CUDA/NCCL workspace, CUDA graph buffers, token imbalance, or dispatcher workspace, so validate final configs with runtime memory metrics.

Quick Decision

When a training run OOMs or is close to the memory limit:

1. Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True first. This fixes fragmentation-induced OOM with zero performance cost. Most Slurm launch templates already include it.
2. For LoRA with sequence parallelism, enable input re-gather (LoRA(sequence_parallel_input_regather=True)). This avoids retaining the full gathered LoRA-A input in every eligible layer; it has no effect when SP is disabled.
3. Add selective activation recompute (recompute_modules=[core_attn]) if not already enabled. See @skills/nemo-mbridge-perf-activation-recompute/SKILL.md.
4. Avoid increasing TP as a memory fix — doubling TP dramatically increases NVLink all-reduce volume and often kills throughput (-28% on Llama3 70B).
5. Avoid increasing PP at the cost of DP — halving DP doubles gradient accumulation steps and hurts throughput (~6%).
6. Consider mlp recompute if still OOM. Saves ~3 GB but costs ~16% GPU utilization on large dense models (Llama3 70B).
7. CPU offloading is blocked when PP > 1.

Enablement

Expandable segments (recommended first step)

Set in the job's environment before launching:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

In Slurm scripts this is typically placed alongside other env vars:

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

No model config changes needed. Zero throughput cost.

Parallelism resizing

If the model genuinely does not fit (not fragmentation), adjust parallelism:

Activation recompute

See @skills/nemo-mbridge-perf-activation-recompute/SKILL.md for full details.

PEFT + sequence-parallel input re-gather

For LoRA training with sequence parallelism, eligible column-parallel linear_qkv and linear_fc1 adapters consume a gathered LayerNorm output. Because LoRA-A is trainable, the default path retains that full gathered input until backward for the LoRA-A weight gradient.

Enable input re-gather when constructing the PEFT config:

from megatron.bridge.peft.lora import LoRA

cfg.peft = LoRA(
    # Keep the recipe's existing LoRA settings here.
    sequence_parallel_input_regather=True,
)

With this option, forward still materializes the full input temporarily for the LoRA-A GEMM, but MCore autograd retains only its sequence-local shard. Backward asynchronously gathers the full input again, overlaps the collective with dgrad when possible, computes the LoRA-A weight gradient, and then reuses the temporary communication buffer.

This is a memory-for-communication tradeoff, not conventional activation checkpointing: no LayerNorm, attention, MLP, or LoRA GEMM is rerun. Some throughput degradation is expected, and the benefit grows with the amount of eligible LoRA-A activation retained. The option has no effect when sequence parallelism is disabled.

CPU offloading

cfg.model.cpu_offloading = True

Incompatible with PP > 1. Only usable when pipeline_model_parallel_size = 1.

A Note on VPP

Virtual pipeline parallelism (VPP) is primarily a throughput optimization that reduces pipeline bubble overhead by interleaving smaller model chunks. Its effect on peak memory is minimal — changing VPP does not meaningfully change the total activation, parameter, or optimizer memory on a GPU.

In earlier experiments we incorrectly attributed an OOM fix to VPP tuning (VPP 5→10). The actual fix was PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True which eliminated memory fragmentation. The VPP=10 run actually used slightly more peak memory (60.2 GB vs 58.8 GB) but did not OOM because expandable segments prevented fragmentation.

VPP should be tuned for pipeline bubble reduction (see @docs/parallelisms.md), not as a memory fix.

Compatibility and Constraints

• expandable_segments:True is incompatible with --use-nccl-ub (NCCL user-buffer registration). See Megatron-FSDP docs.
• When using CUDA graphs with expandable_segments:True, set NCCL_GRAPH_REGISTER=0 (required on pre-Blackwell GPUs, enforced by MCore CudaGraphManager).
• CPU offloading requires pipeline_model_parallel_size = 1.
• Distributed optimizer requires use_distributed_optimizer = True in the optimizer config.
• sequence_parallel_input_regather applies only to eligible non-expert column-parallel LoRA-A projections. Row-parallel adapters, expert adapters, TP=1, CUDA graphs, CPU activation offload, and overlapping full-layer or selective MLP activation recompute fall back to the existing path.

Measured Results

Llama3 70B SFT on 32x H100 80GB, FP8 (Current Scaling):

• Baseline: TP=4, PP=4, VPP=5, DP=2, MBS=1, GBS=32, seq_len=4096
• Golden GPU utilization: 709.93 TFLOP/s/GPU
• Regression threshold: 5%