# Nemo Mbridge Perf Memory Tuning

> Techniques for reducing peak GPU memory in Megatron Bridge, including expandable segments, PEFT plus sequence-parallel input re-gather, parallelism resizing, activation recompute, CPU offloading constraints, and common OOM fixes. Use for GPU OOMs, inadequate memory headroom, LoRA or PEFT activation pressure, memory fragmentation, and memory regressions.

- Skill: `nvidia-nemo-megatron-bridge/nemo-mbridge-perf-memory-tuning` (Agent Skill, multi-file: 6 files)
- Install (CLI): `npx skillmds@latest add nvidia-nemo-megatron-bridge/nemo-mbridge-perf-memory-tuning`
- Raw SKILL.md: https://api.skillmd.com/api/skills/nvidia-nemo-megatron-bridge/nemo-mbridge-perf-memory-tuning/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: AI & ML
- License: Apache-2.0
- Author: NVIDIA-NeMo (https://skillmd.com/u/nvidia-nemo-megatron-bridge)
- Updated: 2026-09-10
- Page: https://skillmd.com/skills/nvidia-nemo-megatron-bridge/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:

```bash
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:

```python
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**.

## After The First Configuration Fits

Do not stop tuning at the first non-OOM layout. Run through optimizer-state
initialization and several steady steps. Inspect W&B memory, then confirm every
rank's peak allocated and reserved memory in the runtime logs; rank 0 can miss
the hot EP or PP rank.

If the hot rank has safe headroom, change one variable at a time:

1. Lower TP by one legal step. This increases per-rank dense state, but can
   remove expensive TP communication and improve efficiency when it still fits.
2. Increase MBS. Adjust gradient accumulation so GBS, data order, and optimizer
   boundaries remain fixed; otherwise this is a convergence change, not a pure
   execution comparison.

Accept the change only when end-to-end tokens/s/GPU or step time improves and
loss remains finite with no skipped/NaN iterations. Do not target 100% reported
memory use: keep margin for checkpoint/eval transients, MoE routing imbalance,
dispatcher workspaces, and CUDA-graph private pools.

## Enablement

### Expandable segments (recommended first step)

Set in the job's environment before launching:

```bash
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
```

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

```bash
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:

```python
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

```python
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%

### Strategy comparison: parallelism changes for memory reduction

| Experiment | TP | PP | VPP | DP | TFLOP/s/GPU | vs Golden | Peak Mem (GB) | Result |
|---|---|---|---|---|---|---|---|---|
| Baseline | 4 | 4 | 5 | 2 | ~704 | -0.8% | 58.8 | OOM (fragmentation) |
| More PP | 4 | 8 | 5 | 1 | 668.0 | -5.9% | 53.2 | Borderline perf |
| More TP | 8 | 4 | 5 | 1 | 508.7 | -28.4% | 50.2 | Severe regression |
| Baseline + expandable_segments | 4 | 4 | 5 | 2 | ~704 | -0.8% | ~59 | **Passed** |

Key takeaways:

- **`expandable_segments:True` is the winner.** The baseline OOM was caused by
  memory fragmentation, not insufficient capacity. Setting this env var
  eliminated the OOM with zero throughput cost and no parallelism changes.
- **PP=8 works for memory but loses DP** (2→1), meaning 32 gradient accumulation
  steps per batch, which hurts throughput by ~6%.
- **TP=8 is catastrophic** (-28%) because doubling TP increases all-reduce
  communication volume proportionally across NVLink, and DP=1 means no
  micro-batch overlap.

### CPU offloading: blocked

| Experiment | offload_layers | Result |
|---|---|---|
| Exp 4 | 2 | Incompatible (PP > 1) |
| Exp 5 | 4 | Incompatible (PP > 1) |
| Exp 6 | 6 | Incompatible (PP > 1) |

`ValueError: Currently there is no support for Pipeline parallelism with CPU
offloading.` This approach is blocked for any model using PP > 1.

### Activation recompute: expensive alternative

Selective activation recompute with `mlp` saved ~3 GB peak memory but cost
~16% GPU utilization on this workload. See
@skills/nemo-mbridge-perf-activation-recompute/SKILL.md for full results.

### LoRA + SP input re-gather

Real-checkpoint H100 training with SQuAD showed lower peak memory in all tested
configurations, with workload-dependent throughput cost:

| Model/config | Baseline peak | Input re-gather peak | Memory saved | Throughput change |
|---|---:|---:|---:|---:|
| Qwen3-8B, TP2, seq 8192 | 47.545 GB | 42.814 GB | 4.731 GB (10.0%) | -6.74% |
| Qwen3-30B-A3B, TP4/EP4 | 29.890 GB | 28.321 GB | 1.569 GB (5.2%) | -2.89% |
| GPT-OSS-120B, TP2/EP8 | 52.185 GB | 51.371 GB | 0.814 GB (1.6%) | -0.34% |

All runs had finite losses with zero skipped or NaN iterations. Two-rank BF16
and FP32 checks matched the baseline for outputs, input gradients, LoRA-A and
LoRA-B gradients, and two-microbatch fused `main_grad` accumulation.

## Code Anchors

### LoRA sequence-parallel input re-gather

```text
src/megatron/bridge/peft/lora.py
    LoRA.sequence_parallel_input_regather

src/megatron/bridge/peft/utils.py
    ParallelLinearAdapter._sequence_parallel_input_regather_eligibility()
    ParallelLinearAdapter.forward()
```

### CPU offloading PP incompatibility (MCore)

```1303:1306:3rdparty/Megatron-LM/megatron/core/transformer/transformer_config.py
        if self.cpu_offloading and self.pipeline_model_parallel_size > 1:
            raise ValueError(
                "Currently there is no support for Pipeline parallelism with CPU offloading"
            )
```

### VPP config and layer divisibility validation (MCore)

```1581:1592:3rdparty/Megatron-LM/megatron/core/transformer/transformer_config.py
            if pipeline_parallel_size and self.virtual_pipeline_model_parallel_size is not None:
                num_layers_per_middle_pipeline_rank = num_layers // pipeline_parallel_size
                if (
                    not num_layers_per_middle_pipeline_rank
                    % self.virtual_pipeline_model_parallel_size
                    == 0
                ):
                    raise ValueError(
                        f"number of layers on each middle pipeline rank:"
                        f"{num_layers_per_middle_pipeline_rank} must be divisible by virtual"
                        f"pipeline parallel degree {self.virtual_pipeline_model_parallel_size}"
                    )
```

### Parallelism docs on interleaved pipeline schedule

```116:124:docs/parallelisms.md
To minimize the pipeline bubble, the computation on each GPU can be divided into multiple subsets of layers (referred to as model chunks), rather than a single contiguous block. Enable this by setting `virtual_pipeline_model_parallel_size`:

model_config = GPTModelProvider(
    pipeline_model_parallel_size=4,
    virtual_pipeline_model_parallel_size=2,  # 2 model chunks per pipeline stage
    # ... other model parameters
)
```

## Failure Diagnosis

| Symptom | Cause | Confirm | Fix |
|---|---|---|---|
| OOM on a single rank despite headroom on others | Memory fragmentation | check if `expandable_segments:True` is set | set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` |
| OOM with `expandable_segments` already set | Genuine capacity limit | check `nvidia-smi` for param/optimizer memory | increase PP, use distributed optimizer, or add recompute |
| Estimated memory exceeds GPU capacity before launch | model state or activations genuinely too large | run `estimate_training_memory` and inspect the largest component | adjust PP/TP/CP/EP, distributed optimizer, or recompute before launching |
| LoRA + SP retains unexpectedly high activation memory | full gathered LoRA-A inputs are retained until backward | check whether `cfg.peft.sequence_parallel_input_regather` is enabled and the target is eligible | set `LoRA(sequence_parallel_input_regather=True)`; verify fallback constraints |
| `ValueError: PP + CPU offloading` | using cpu_offloading with PP > 1 | check PP config | disable CPU offloading or set PP=1 |
| `RuntimeError` with `--use-nccl-ub` + expandable segments | NCCL UB incompatible with expandable allocator | check env vars | remove `expandable_segments:True` or disable `--use-nccl-ub` |

## Known Limitations

- CPU offloading is blocked when PP > 1
- Parallelism resizing (TP/PP) often has significant throughput costs
- The theoretical estimator is formula-based and does not replace runtime
  profiling or CUDA memory reports
- LoRA input re-gather does not cover row-parallel or expert adapters and may
  have negligible benefit when few eligible LoRA-A activations dominate memory
- W&B or logger memory is an observation surface, not an automatic optimizer;
  TP and MBS candidates still require matched steady-state validation

## Verification

Quick check that `expandable_segments:True` is active:

```python
import os
assert "expandable_segments:True" in os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
```

For Slurm jobs, verify the env var is exported before the training command
in the launch script.

For LoRA + SP input re-gather, run the focused configuration tests and the real
two-rank MCore backward-parity test:

```bash
uv run python -m pytest \
  tests/unit_tests/peft/test_utils.py -k "sequence_parallel_input_regather" \
  tests/unit_tests/peft/test_lora.py -k "sequence_parallel_input_regather"

uv run python -m torch.distributed.run --nproc_per_node=2 -m pytest \
  tests/unit_tests/peft/test_lora_sp_input_regather_distributed.py
```