/nemo-mbridge-perf-memory-tuning

name: nemo-mbridge-perf-memory-tuning
description: Techniques for reducing peak GPU memory in Megatron Bridge — expandable segments, PEFT + SP input re-gather, parallelism resizing, activation recompute, CPU offloading constraints, and common OOM fixes.
license: Apache-2.0
when_to_use: GPU OOM errors, reducing peak memory, reducing LoRA or PEFT activation memory with sequence parallelism, or tracing an OOM regression to a specific commit or config change; 'out of memory', 'OOM', 'memory fragmentation', 'expandable_segments', 'reduce GPU memory', 'LoRA memory', 'PEFT memory', 'sequence_parallel_input_regather', 'PYTORCH_CUDA_ALLOC_CONF'.

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:

| Strategy | Memory effect | Throughput cost | Notes | |---|---|---|---| | Increase PP (keeping DP) | Fewer layers per stage | Moderate (~6% if DP halved) | Only if GPU count allows | | Increase TP | Fewer params per GPU | Severe (-28% on 70B) | Last resort | | Distributed optimizer | Shards optimizer state across DP ranks | ~1-2% | Recommended for large models | | FSDP | Shards params + grads + optimizer | Varies | See @skills/nemo-mbridge-perf-megatron-fsdp/SKILL.md |

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 ove