NeMo MBridge GPU Memory Tuning
Diagnose GPU OOMs in Megatron Bridge training and reduce fragmentation, activation, and PEFT memory usage.
The skill recommends environment variables and training configuration changes only; no malware, credential collection, covert exfiltration, or destructive defaults are evident. It discloses NCCL, CUDA Graph, PP, and LoRA constraints and gives rollback directions such as removing an incompatible setting. Six points are deducted for lacking explicit user confirmation, pre-change backup/rollback procedure, dependency versions, and fuller data-flow boundaries. NVIDIA’s verified provenance supports attribution and governance only, not intrinsic safety.
The document provides a decision sequence, compatibility constraints, symptom/confirmation/fix tables, and coherent key paths on static inspection. Eleven points are deducted because execution depends on Megatron Bridge, PyTorch, MCore, CUDA, NCCL, and workload-specific configuration; no commands or tests were executed here, and abnormal-input coverage and reproducibility remain limited. The score respects the static-review cap.
Triggers, audience, and OOM or LoRA-plus-sequence-parallel scenarios are specific, with boundaries for TP, PP, VPP, CPU offloading, and adapter eligibility. Core use does not depend on an overseas online service, so no mainland-China reachability deduction applies. Four points are deducted because Chinese-language support is absent and hardware, software-version, model, and cluster compatibility ranges are incompletely specified.
The skill has clear sections for quick decisions, configuration examples, constraints, measured results, diagnosis, limitations, and verification commands; the card adds owner, version, evaluation, and Apache-2.0 licensing information. Five points are deducted for missing recommended Instructions and Examples sections, absent author metadata, incomplete installation/dependency notes, and no explicit changelog or concrete maintenance responsibility.
For the stated Megatron Bridge OOM cases, it gives directly usable first-step configuration, alternatives, performance tradeoffs, and compatibility warnings. The supplied Llama3, LoRA, and CPU-offloading cases support practical value. Eight points are deducted because the output and fixes were not executed in this review, coverage relies on referenced external files, and there is no automatic profiling or environment-specific strategy selection. The score respects the static cap.
Code anchors, test commands, failure messages, measurement tables, PR identifiers, and numerical/gradient parity claims provide auditable primary material. Five points are deducted because tests were not executed, the cited PR and measurements cannot be independently reproduced from the supplied material, and corroboration across independent sources is thin. The score respects the static cap.
- Do not treat expandable_segments as a universal capacity fix; confirm fragmentation first and account for its incompatibility with --use-nccl-ub and some CUDA Graph configurations.
- Before production training, obtain user confirmation for configuration changes and validate against the actual GPU, CUDA, NCCL, and Megatron Bridge versions and runtime metrics; CPU offloading is unavailable with PP>1, and LoRA input re-gather adds communication and may reduce throughput.
- The benchmark contains only one positive task, and this review executed no tests; broad effectiveness across other models, hardware, or clusters is unverified.
What does this skill do, and when should you use it?
This skill targets GPU out-of-memory failures in NeMo MBridge and Megatron Bridge training jobs. It prioritizes CUDA allocator fragmentation and also covers LoRA with sequence parallelism, activation recompute, parallelism resizing, distributed optimizers, FSDP, and CPU offloading. It provides a theoretical memory estimator workflow, compatibility constraints, failure-diagnosis guidance, and verification tests. Its measured results focus on 32 H100 80GB GPUs running Llama3 70B and several LoRA sequence-parallel configurations.
It instructs users to set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, call Megatron Bridge's estimate_training_memory estimator, and analyze parameter, optimizer, activation, and temporary workspace memory. It recommends changes involving TP, PP, DP, distributed optimizers, FSDP, and activation recompute; shows how to configure LoRA(sequence_parallel_input_regather=True) to reduce retained gathered LoRA-A inputs; and checks constraints involving CPU offloading, CUDA graphs, NCCL user buffers, and adapter types. It also supplies environment checks, pytest commands, and a two-process distributed parity test.
- A Megatron Bridge job OOMs on one rank while other GPUs still have headroom, suggesting allocator fragmentation.
- A LoRA or PEFT job using sequence parallelism retains excessive activation memory.
- A model genuinely exceeds GPU capacity and the user needs to compare PP, TP, distributed optimizer, FSDP, or recompute options.
- A configuration change or commit appears to have introduced an OOM regression and needs targeted diagnosis.
- A user wants to enable CPU offloading or CUDA graphs and needs to validate pipeline-parallel and CUDA/NCCL constraints.
What are this skill's strengths and limitations?
- Starts with expandable segments, which the supplied benchmark shows can remove fragmentation-induced OOM without changing the model or parallelism.
- Covers LoRA sequence-parallel input re-gather with measured memory savings and throughput changes across several configurations.
- Clearly documents tradeoffs and blockers involving TP, PP, VPP, CPU offloading, CUDA graphs, and NCCL user buffers.
- Includes estimator usage, failure diagnosis, code anchors, and copyable verification commands.
- The measurements are concentrated on H100 80GB hardware and selected Llama3, Qwen3, and GPT-OSS workloads, so results may not generalize.
- The theoretical estimator excludes allocator fragmentation, CUDA/NCCL workspace, CUDA graph buffers, token imbalance, and dispatcher workspace.
- Activation recompute and parallelism changes can materially reduce throughput; the TP=8 example lost 28.4%.
- CPU offloading is unavailable when pipeline parallelism exceeds one, and input re-gather does not cover every LoRA adapter.
How do you install this skill?
Install the specific skill with the README-supported Skills CLI command: npx skills add nvidia/skills --skill nemo-mbridge-perf-memory-tuning --yes. The README does not document additional skill-specific installation steps; after installation, the skill is available when the agent next loads skills and encounters a relevant task.
How do you use this skill?
Give the agent a concrete training problem, for example: “My Megatron Bridge training job has a GPU OOM. Check for allocator fragmentation first, then evaluate LoRA sequence-parallel input re-gather and activation recompute.” For verification, use the supplied commands to export the expandable-segments setting, run estimate_training_memory, and execute the focused pytest and two-process torch.distributed tests.
How does this skill compare with similar options?
Compared with immediately increasing TP or PP, the skill prioritizes fixing allocator fragmentation; the supplied results show nearly no throughput cost for expandable segments, versus a 28.4% throughput loss for TP=8 and about 5.9% for PP=8 after reducing DP. Compared with activation recompute, LoRA input re-gather trades memory for communication and does not rerun LayerNorm, attention, MLP, or LoRA GEMMs. VPP is presented primarily as a pipeline-bubble optimization, not a memory fix.