{"id":"6c6fb7a9-a5be-4406-be65-e512c3e31cf9","slug":"nemo-mbridge-perf-memory-tuning","kind":"skill","name":"nemo-mbridge-perf-memory-tuning","summary":"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 fixe","owner":"NVIDIA","repoUrl":"https://github.com/NVIDIA/skills/tree/HEAD/skills/nemo-mbridge-perf-memory-tuning","license":"Apache-2.0","licenseTier":"permissive","trust":"official","stars":3586,"match":null,"fewStars":false,"format":"skill","createdAt":null,"addedAt":"2026-10-14T05:47:37.956Z","logoUrl":"https://avatars.githubusercontent.com/u/1728152?v=4","linkOnly":false,"collection":{"name":"skills","repo":"nvidia/skills"},"skillPath":"skills/nemo-mbridge-perf-memory-tuning","categories":["memory-knowledge","developer-tools"],"tags":["agentic-ai","ai-agents","claude-code","codex","coding-agents","cuda","developer-tools","llm"],"runtime":"binary","local":true,"cloudCapable":false,"cloudCapability":"unknown","cloudReadyOnPlatform":false,"needsApiKey":false,"reviewed":false,"syncedAt":"2026-10-14T05:47:37.955Z","sample":false,"url":"https://conduits.dev/tools/nemo-mbridge-perf-memory-tuning","language":"Python","transports":[],"install":{},"requirements":{"env":[],"upstreamMayCharge":false},"tools":null,"readme":"> 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.\n\n# Memory Tuning\n\nStable docs: @docs/parallelisms.md\nCard: @skills/nemo-mbridge-perf-memory-tuning/card.yaml\n\n## What It Is\n\nGPU OOM failures during training often stem from memory **fragmentation** rather\nthan raw capacity.  PyTorch's default CUDA allocator can leave unusable gaps\nbetween allocations.  The single most effective fix is:\n\n```bash\nexport PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True\n```\n\nThis tells PyTorch to use expandable (non-fixed-size) memory segments, which\ndramatically reduces fragmentation and often eliminates borderline OOM without\nany model or parallelism changes.\n\nBeyond fragmentation, actual peak memory is determined by:\n\n- **Parameter + optimizer state memory** — controlled by TP, PP, DP sharding\n  (distributed optimizer, FSDP)\n- **Activation memory** — controlled by activation recompute, sequence length,\n  micro-batch size, and PEFT-specific retention of gathered inputs\n- **Temporary / workspace memory** — CUDA kernels, NCCL buffers, CUDA graphs\n\nFor configuration planning, use the Bridge theoretical estimator before launching\nlarge jobs:\n\n```python\nfrom megatron.bridge.training.utils.theoretical_memory_utils import estimate_training_memory\n\nestimate = estimate_training_memory(cfg, num_microbatches=num_microbatches)\n```\n\nThe estimator reports the most-loaded GPU shard and separates dense/embedding,\nrouted MoE expert, and activation components. It does not include allocator\nfragmentation, CUDA/NCCL workspace, CUDA graph buffers, token imbalance, or\ndispatcher workspace, so validate final configs with runtime memory metrics.\n\n## Quick Decision\n\nWhen a training run OOMs or is close to the memory limit:\n\n1. **Set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` first.** This fixes\n   fragmentation-induced OOM with zero performance cost. Most Slurm launch\n   templates already include it.\n2. **For LoRA with sequence parallelism, enable input re-gather**\n   (`LoRA(sequence_parallel_input_regather=True)`). This avoids retaining the\n   full gathered LoRA-A input in every eligible layer; it has no effect when SP\n   is disabled.\n3. **Add selective activation recompute** (`recompute_modules=[core_attn]`) if\n   not already enabled. See @skills/nemo-mbridge-perf-activation-recompute/SKILL.md.\n4. **Avoid increasing TP** as a memory fix — doubling TP dramatically increases\n   NVLink all-reduce volume and often kills throughput (-28% on Llama3 70B).\n5. **Avoid increasing PP at the cost of DP** — halving DP doubles gradient\n   accumulation steps and hurts throughput (~6%).\n6. Consider `mlp` recompute if still OOM. Saves ~3 GB but costs ~16% GPU\n   utilization on large dense models (Llama3 70B).\n7. CPU offloading is **blocked when PP > 1**.\n\n## Enablement\n\n### Expandable segments (recommended first step)\n\nSet in the job's environment before launching:\n\n```bash\nexport PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True\n```\n\nIn Slurm scripts this is typically placed alongside other env vars:\n\n```bash\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\nexport NVTE_ALLOW_NONDETERMINISTIC_ALGO=1\nexport PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True\n```\n\nNo model config changes needed. Zero throughput cost.\n\n### Parallelism resizing\n\nIf the model genuinely does not fit (not fragmentation), adjust parallelism:\n\n| Strategy | Memory effect | Throughput cost | Notes |\n|---|---|---|---|\n| Increase PP (keeping DP) | Fewer layers per stage | Moderate (~6% if DP halved) | Only if GPU count allows |\n| Increase TP | Fewer params per GPU | Severe (-28% on 70B) | Last resort |\n| Distributed optimizer | Shards optimizer state across DP ranks | ~1-2% | Recommended for large models |\n| FSDP | Shards params + grads + optimizer | Varies | See @skills/nemo-mbridge-perf-megatron-fsdp/SKILL.md |\n\n### Activation recompute\n\nSee @skills/nemo-mbridge-perf-activation-recompute/SKILL.md for full details.\n\n### PEFT + sequence-parallel input re-gather\n\nFor `LoRA` training with sequence parallelism, eligible column-parallel\n`linear_qkv` and `linear_fc1` adapters consume a gathered LayerNorm output.\nBecause LoRA-A is trainable, the default path retains that full gathered input\nuntil backward for the LoRA-A weight gradient.\n\nEnable input re-gather when constructing the PEFT config:\n\n```python\nfrom megatron.bridge.peft.lora import LoRA\n\ncfg.peft = LoRA(\n    # Keep the recipe's existing LoRA settings here.\n    sequence_parallel_input_regather=True,\n)\n```\n\nWith this option, forward still materializes the full input temporarily for the\nLoRA-A GEMM, but MCore autograd retains only its sequence-local shard. Backward\nasynchronously gathers the full input again, overlaps the collective with\ndgrad when possible, computes the LoRA-A weight gradient, and then reuses the\ntemporary communication buffer.\n\nThis is a memory-for-communication tradeoff, not conventional activation\ncheckpointing: no LayerNorm, attention, MLP, or LoRA GEMM is rerun. Some\nthroughput degradation is expected, and the benefit grows with the amount of\neligible LoRA-A activation retained. The option has no effect when sequence\nparallelism is disabled.\n\n### CPU offloading\n\n```python\ncfg.model.cpu_offloading = True\n```\n\n**Incompatible with PP > 1.** Only usable when `pipeline_model_parallel_size = 1`.\n\n## A Note on VPP\n\nVirtual pipeline parallelism (VPP) is primarily a **throughput** optimization\nthat reduces pipeline bubble overhead by interleaving smaller model chunks. Its\neffect on peak memory is minimal — changing VPP does not meaningfully change\nthe total activation, parameter, or optimizer memory on a GPU.\n\nIn earlier experiments we incorrectly attributed an OOM fix to VPP tuning\n(VPP 5→10). The actual fix was `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`\nwhich eliminated memory fragmentation. The VPP=10 run actually used slightly\n**more** peak memory (60.2 GB vs 58.8 GB) but did not OOM because expandable\nsegments prevented fragmentation.\n\nVPP should be tuned for pipeline bubble reduction (see @docs/parallelisms.md),\nnot as a memory fix.\n\n## Compatibility and Constraints\n\n- `expandable_segments:True` is incompatible with `--use-nccl-ub` (NCCL\n  user-buffer registration). See Megatron-FSDP docs.\n- When using CUDA graphs with `expandable_segments:True`, set\n  `NCCL_GRAPH_REGISTER=0` (required on pre-Blackwell GPUs, enforced by MCore\n  `CudaGraphManager`).\n- CPU offloading requires `pipeline_model_parallel_size = 1`.\n- Distributed optimizer requires `use_distributed_optimizer = True` in the\n  optimizer config.\n- `sequence_parallel_input_regather` applies only to eligible non-expert\n  column-parallel LoRA-A projections. Row-parallel adapters, expert adapters,\n  TP=1, CUDA graphs, CPU activation offload, and overlapping full-layer or\n  selective MLP activation recompute fall back to the existing path.\n\n## Measured Results\n\nLlama3 70B SFT on 32x H100 80GB, FP8 (Current Scaling):\n- Baseline: TP=4, PP=4, VPP=5, DP=2, MBS=1, GBS=32, seq_len=4096\n- Golden GPU utilization: 709.93 TFLOP/s/GPU\n- Regression threshold: 5%\n\n### Strategy comparison: parallelism changes for memory reduction\n\n| Experiment | TP | PP | VPP | DP | TFLOP/s/GPU | vs Golden | Peak Mem (GB) | Result |\n|---|---|---|---|---|---|---|---|---|\n| Baseline | 4 | 4 | 5 | 2 | ~704 | -0.8% | 58.8 | OOM (fragmentation) |\n| More PP | 4 | 8 | 5 | 1 | 668.0 | -5.9% | 53.2 | Borderline perf |\n| More TP | 8 | 4 | 5 | 1 | 508.7 | -28.4% | 50.2 | Severe regression |\n| Baseline + expandable_segments | 4 | 4 | 5 | 2 | ~704 | -0.8% | ~59 | **Passed** |\n\nKey takeaways:\n\n- **`expandable_segments:True` is the winner.** The baseline OOM was caused by\n  memory fragmentation, not insufficient capacity. Setting this env var\n  eliminated the OOM with zero throughput cost and no parallelism changes.\n- **PP=8 works for memory but loses DP** (2→1), meaning 32 gradient accumulation\n  steps per batch, which hurts throughput by ~6%.\n- **TP=8 is catastrophic** (-28%) because doubling TP increases all-reduce\n  communication volume proportionally across NVLink, and DP=1 means no\n  micro-batch overlap.\n\n### CPU offloading: blocked\n\n| Experiment | offload_layers | Result |\n|---|---|---|\n| Exp 4 | 2 | Incompatible (PP > 1) |\n| Exp 5 | 4 | Incompatible (PP > 1) |\n| Exp 6 | 6 | Incompatible (PP > 1) |\n\n`ValueError: Currently there is no support for Pipeline parallelism with CPU\noffloading.` This approach is blocked for any model using PP > 1.\n\n### Activation recompute: expensive alternative\n\nSelective activation recompute with `mlp` saved ~3 GB peak memory but cost\n~16% GPU utilization on this workload. See\n@skills/nemo-mbridge-perf-activation-recompute/SKILL.md for full results.\n\n### LoRA + SP input re-gather\n\nReal-checkpoint H100 training with SQuAD showed lower peak memory in all tested\nconfigurations, with workload-dependent throughput cost:\n\n| Model/config | Baseline peak | Input re-gather peak | Memory saved | Throughput change |\n|---|---:|---:|---:|---:|\n| Qwen3-8B, TP2, seq 8192 | 47.545 GB | 42.814 GB | 4.731 GB (10.0%) | -6.74% |\n| Qwen3-30B-A3B, TP4/EP4 | 29.890 GB | 28.321 GB | 1.569 GB (5.2%) | -2.89% |\n| GPT-OSS-120B, TP2/EP8 | 52.185 GB | 51.371 GB | 0.814 GB (1.6%) | -0.34% |\n\nAll runs had finite losses with zero skipped or NaN iterations. Two-rank BF16\nand FP32 checks matched the baseline for outputs, input gradients, LoRA-A and\nLoRA-B gradients, and two-microbatch fused `main_grad` accumulation.\n\n## Code Anchors\n\n### LoRA sequence-parallel input re-gather\n\n```text\nsrc/megatron/bridge/peft/lora.py\n    LoRA.sequence_parallel_input_regather\n\nsrc/megatron/bridge/peft/utils.py\n    ParallelLinearAdapter._sequence_parallel_input_regather_eligibility()\n    ParallelLinearAdapter.forward()\n```\n\n### CPU offloading PP incompatibility (MCore)\n\n```1303:1306:3rdparty/Megatron-LM/megatron/core/transformer/transformer_config.py\n        if self.cpu_offloading and self.pipeline_model_parallel_size > 1:\n            raise ValueError(\n                \"Currently there is no support for Pipeline parallelism with CPU offloading\"\n            )\n```\n\n### VPP config and layer divisibility validation (MCore)\n\n```1581:1592:3rdparty/Megatron-LM/megatron/core/transformer/transformer_config.py\n            if pipeline_parallel_size and self.virtual_pipeline_model_parallel_size is not None:\n                num_layers_per_middle_pipeline_rank = num_layers // pipeline_parallel_size\n                if (\n                    not num_layers_per_middle_pipeline_rank\n                    % self.virtual_pipeline_model_parallel_size\n                    == 0\n                ):\n                    raise ValueError(\n                        f\"number of layers on each middle pipeline rank:\"\n                        f\"{num_layers_per_middle_pipeline_rank} must be divisible by virtual\"\n                        f\"pipeline parallel degree {self.virtual_pipeline_model_parallel_size}\"\n                    )\n```\n\n### Parallelism docs on interleaved pipeline schedule\n\n```116:124:docs/parallelisms.md\nTo 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`:\n\nmodel_config = GPTModelProvider(\n    pipeline_model_parallel_size=4,\n    virtual_pipeline_model_parallel_size=2,  # 2 model chunks per pipeline stage\n    # ... other model parameters\n)\n```\n\n## Failure Diagnosis\n\n| Symptom | Cause | Confirm | Fix |\n|---|---|---|---|\n| 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` |\n| OOM with `expandable_segments` already set | Genuine capacity limit | check `nvidia-smi` for param/optimizer memory | increase PP, use distributed optimizer, or add recompute |\n| 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 |\n| 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 |\n| `ValueError: PP + CPU offloading` | using cpu_offloading with PP > 1 | check PP config | disable CPU offloading or set PP=1 |\n| `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` |\n\n## Known Limitations\n\n- CPU offloading is blocked when PP > 1\n- Parallelism resizing (TP/PP) often has significant throughput costs\n- The theoretical estimator is formula-based and does not replace runtime\n  profiling or CUDA memory reports\n- LoRA input re-gather does not cover row-parallel or expert adapters and may\n  have negligible benefit when few eligible LoRA-A activations dominate memory\n\n## Verification\n\nQuick check that `expandable_segments:True` is active:\n\n```python\nimport os\nassert \"expandable_segments:True\" in os.environ.get(\"PYTORCH_CUDA_ALLOC_CONF\", \"\")\n```\n\nFor Slurm jobs, verify the env var is exported before the training command\nin the launch script.\n\nFor LoRA + SP input re-gather, run the focused configuration tests and the real\ntwo-rank MCore backward-parity test:\n\n```bash\nuv run python -m pytest \\\n  tests/unit_tests/peft/test_utils.py -k \"sequence_parallel_input_regather\" \\\n  tests/unit_tests/peft/test_lora.py -k \"sequence_parallel_input_regather\"\n\nuv run python -m torch.distributed.run --nproc_per_node=2 -m pytest \\\n  tests/unit_tests/peft/test_lora_sp_input_regather_distributed.py\n```"}