Agent skill · NVIDIA
nemo-mbridge-perf-memory-tuning
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.
What it needs
About 7k tokens when loaded.
What this skill does
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: 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: 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 PYTORCHCUDAALLOCCONF=expandablesegments: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(sequenceparallelinputregather=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 (recomputemodules=[coreattn]) if not already enabled. See @skills/nemo-mbridge-perf-activation-recompute/SKILL.md. 4. …
How to use it
Reference it in AdaL, Claude Code, Cursor or any coding agent — nothing to install:
@skills NVIDIA/nemo-mbridge-perf-memory-tuning