复制安装命令
用 Codex 或 Claude 安装复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它先审查 Skill 页面再帮你安装。
复制前请先查看来源、License 和安全提示。
SkillAtlasSkill 详情
nemo-mbridge-perf-memory-tuning
Official, NVIDIA-verified Agent Skills for Claude Code, Codex, and other coding agents.
审核状态：已审核Quality 72Security 70
项目 README
来源文件：README.md
NVIDIA Agent Skills
Official, NVIDIA-verified Agent Skills for Claude Code, Codex, and other coding agents.
📖 Docs: docs.nvidia.com/skills · 📺 Livestream: From Vulnerable to Verified · 📝 Blog: NVIDIA Verified Agent Skills: Capability Governance for AI Agents
Skills are portable instruction sets that teach AI agents how to use NVIDIA software optimally: Physical AI and robotics workflows, simulation, CUDA-X libraries, RAG and AI Blueprints, and platform tools. This repository is a catalog: skills are maintained in their respective product repos, and mirrored here daily via an automated sync pipeline. Skills are being added continuously, so check back for updates. We are building this infrastructure in the open, and contributions are welcome. See the Roadmap for what is planned next.
Quickstart
Install NVIDIA skills with the default skills CLI flow:
npx skills add nvidia/skills
The CLI runs through npx and prompts you to choose a skill and install destination. You do not need to clone this repo or copy skill folders by hand.
Requires a current
skillsCLI (v1.5.16 or newer). Installing vianpx skills@latest add nvidia/skillsalways uses the latest. On older CLIs (v1.5.15 and earlier), skills may install but not appear in Claude Code — see Troubleshooting.
The skill is available the next time your agent loads skills and encounters a relevant task. For example, ask your agent to "solve a linear programming problem with cuOpt" and the skill guides it through the cuOpt Python API. In Claude Code, run /reload-skills to load newly installed skills in your current session.
Install One Skill Without Prompts
Use this when you already know the skill name and want to skip prompts.
npx skills add nvidia/skills --skill cuopt-numerical-optimization-api --yes
Replace cuopt-numerical-optimization-api with any skill name from the Skill Catalog.
Install for a Specific Agent
Use --agent to target a specific AI coding agent. Initially, we'll support common client targets, expanding the list over time. For the full list of clients supported by the spec, see the skills CLI Supported Agents table.
Claude Code
npx skills add nvidia/skills --skill cuopt-numerical-optimization-api --agent claude-code
Codex
npx skills add nvidia/skills --skill cuopt-numerical-optimization-api --agent codex
Snowflake CoCo
npx skills add nvidia/skills --skill cuopt-numerical-optimization-api --agent cortex
Cursor
npx skills add nvidia/skills --skill cuopt-numerical-optimization-api --agent cursor
Kiro
npx skills add nvidia/skills --skill cuopt-numerical-optimization-api --agent kiro-cli
Use --agent more than once to install the same skill into multiple agents.
npx skills add nvidia/skills \
--skill cuopt-numerical-optimization-api \
--agent claude-code \
--agent codex \
--agent cursor \
--agent kiro-cli
Keep Skills Up to Date
New skills land continuously, and existing ones are revised, renamed, or consolidated as the catalog evolves. Refresh what you have installed with:
npx skills update
Run it interactively and the CLI also flags skills that were removed or merged upstream (for example, when several skills are consolidated into one) and offers to remove the stale local copies. Use npx skills list to see what is installed and npx skills check to preview what is out of date first.
Browse the Catalog
Use this when you want to see available NVIDIA skills before installing anything.
npx skills add nvidia/skills --list
For non-interactive installs, global installs, agent-specific installs, updates, removals, and fallback manual copying, see Advanced installation.
Skill Catalog
Getting Help & Contributing
Where to file an issue depends on what's broken:
- Skill content issues (a specific skill has a bug, missing functionality, or incorrect content) — file in the source repo for that product, using the per-product table below.
- Catalog issues (catalog README errors, sync workflow problems, distribution channels, signing/verification flow, docs in this repo) — file here using the catalog issue templates: Bug Report, Feature Request, or Documentation Request or Correction.
- Questions or general discussion — use Discussions. The issue tracker is reserved for bug reports, feature proposals with a design, and documentation issues.
- Security vulnerabilities — follow the disclosure process in SECURITY.md; do not open a public issue.
Per-product source repo links:
For issues with this catalog repo itself (README, structure, listing a new product): open an issue here.
Verifying Skills
Every published skill ships with a detached OMS signature (skill.oms.sig). The sync pipeline drops any skill missing the required artifacts before publishing, so every skill in the catalog carries:
SKILL.md— the skill instructions consumed by the agentskill-card.md— skill identity and governance cardskill.oms.sig— detached OMS signature (verifiable againstnv-agent-root-cert.pem)- A Tier-3 evaluation dataset — accepted at
evals/evals.json,evals/*.json,eval/*.json, orbenchmark/evals.json BENCHMARK.md— generated benchmark report capturing verifiable uplift data
Verify a skill against the NVIDIA trust anchor nv-agent-root-cert.pem:
pip install model-signing
model_signing verify certificate SKILL_DIR \
--signature SKILL_DIR/skill.oms.sig \
--certificate_chain nv-agent-root-cert.pem \
--ignore_unsigned_files
A successful verification confirms that the skill contents have not been modified since signing by NVIDIA.
See Verify Signed Agent Skills for signature layout, the trust pipeline, and policy options.
Roadmap
- ✅ Public skills catalog with NVIDIA-verified skills across multiple products
- ✅ Automated sync pipeline with skills mirrored from product repos daily
- ✅ Security scanning for all published skills covering instruction safety and supply-chain integrity
- ✅ Skills signing so every published skill carries a verifiable NVIDIA signature
- ✅ Skills universal evaluation criteria and task-specific criteria
- ✅ Skill Card with machine-readable metadata for identity, provenance, quality, and behavioral boundaries
- ✅ Sync-time compliance gates — signature drift detection and missing-artifact enforcement
- ✅ Syndication to external marketplaces — Skills.sh, Codex plugin, Claude Code plugin, ClawHub, Hermes Hub
- 🔲 Syndication to additional MCP hubs and partner channels
Repository Structure
NVIDIA/skills/
├── skills/ # NVIDIA-verified skills (count grows continuously),
│ │ synced from upstream product repos
│ ├── README.md # Browser-facing install guidance
│ ├── <product-prefix>-*/ # Flat layout — one dir per skill, product-prefixed
│ │ # e.g. aiq-*, cuopt-*, cupynumeric-*,
│ │ # dali-*, deepstream-*, dicom-*, digital-health-*,
│ │ # dynamo-*, earth2studio-*, holoscan-*, hsb-*,
│ │ # jetson-*, launch-nemo-rl, mcore-*,
│ │ # nemo-automodel-*, nemo-data-designer-plugin,
│ │ # nemo-evaluator-plugin, nemo-mbridge-* (20 skills),
│ │ # nemo-retriever, nemo-rl-* (4 skills),
│ │ # nemoclaw-user-guide, nemotron-*, nemotron-speech,
│ │ # nv-* (medical AI), physicsnemo-*, rag-*,
│ │ # skill-card-generator, tao-*, tilegym-*,
│ │ # vss-* (15 skills), accelerated-computing-cudf,
│ │ # cudaq-guide, portfolio-optimization
│ ├── omniverse-*/ # Physical AI — manually staged (see manual-components.yml)
│ └── physical-ai-*/ # Physical AI — manually staged
├── components.d/ # Product registry — one file per component, teams onboard here
│ ├── README.md # Schema and onboarding instructions
│ └── <product>.yml # one file per registered product
├── plugins/ # Packaged plugin distributions
│ └── nvidia-skills/ # Curated NVIDIA skills bundle (Claude Code, Codex)
├── plugins.d/ # Plugin build registry — config for `build-plugins.py`
│ ├── README.md
│ ├── _defaults.yml
│ └── nvidia-skills.yml
├── .claude-plugin/ # Claude Code marketplace metadata
│ └── marketplace.json
├── .agents/plugins/ # Agent marketplace metadata (other clients)
│ └── marketplace.json
├── docs/ # Long-form documentation (published via Fern)
│ ├── README.md # How to build the docs locally
│ ├── index.mdx
│ ├── advanced-install.mdx
│ ├── agent-skill-trust-pipeline.mdx
│ ├── release-checklist.mdx
│ ├── scanning-agent-skills.mdx
│ ├── signing-agent-skills.mdx
│ └── skill-cards.mdx
├── fern/ # Fern docs site configuration
├── .github/
│ ├── workflows/ # Sync pipeline, plugin validation, DCO check, author verify
│ └── scripts/ # regenerate-readme.sh, build-plugins.py,
│ # manual-components.yml (temp Physical AI catalog
│ # exception, removed after Computex 2026),
│ # marketplace/metadata.json (skill metadata sidecar)
├── nv-agent-root-cert.pem # Trust anchor for OMS signature verification
├── skills.sh.json # Skills.sh marketplace grouping config
├── CHANGELOG.md
├── CONTRIBUTING.md # Contribution guidelines
├── SECURITY.md # Security reporting policy
├── CODE_OF_CONDUCT.md # Community code of conduct
├── LICENSE-APACHE # Apache 2.0 (source code)
└── LICENSE-CC-BY-4.0 # CC BY 4.0 (documentation/skills)
Skills are maintained in their respective product repos (see the Source column in the Skill Catalog) and synced to this repo daily. Products only appear under skills/ after the sync pipeline confirms each skill carries:
skill.oms.sig— detached OMS-format signature (verifiable againstnv-agent-root-cert.pem)skill-card.md— skill identity and governance card- A Tier-3 evaluation dataset — accepted at
evals/evals.json,evals/*.json,eval/*.json, orbenchmark/evals.json
When evaluation runs produce a BENCHMARK.md, it ships alongside the skill so consumers can see verifiable benchmark uplift data.
Standards & Compatibility
This repository adheres to the Agent Skills specification:
- Skills are portable directories with a
SKILL.mdfile at their root. - Metadata uses YAML frontmatter with required
nameanddescriptionfields. - Skills follow a progressive disclosure model — lightweight metadata loads at startup, full instructions load on activation.
- Validate your skill using the
skills-refreference library.
License
Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
This code is dual-licensed with documentation/skills under the CC-BY-4.0 AND source code under Apache-2.0 license terms. The full license texts can be found in LICENSE-APACHE and LICENSE-CC-BY-4.0 respectively.
其他
中风险
- 来源需自行核对维护者身份。
- 包含脚本或命令调用，安装前请复核。
- 可能需要外部 token、网络权限或第三方服务。
- 未检测到高风险命令。
- 扫描发现：2 条。
Codex — Git Clone 安装
- 安装前请先查看来源仓库和风险报告。
- 克隆仓库：git clone https://github.com/NVIDIA/skills.git
- 将 "skills/nemo-mbridge-perf-memory-tuning" 文件夹复制到 Codex 的 skills 目录中。
- 重启 Codex 让新的 skill 生效。
Codex — 手动复制安装
- 安装前请先查看来源仓库和风险报告。
- 从源仓库下载 SKILL.md 及相关文件。
- 在 Codex 的 skills 目录中创建新文件夹。
- 将所有 skill 文件复制到新文件夹中。
- 重启 Codex 让新的 skill 生效。
Claude Code — Git Clone 安装
- 安装前请先查看来源仓库和风险报告。
- 克隆仓库：git clone https://github.com/NVIDIA/skills.git
- 将 "skills/nemo-mbridge-perf-memory-tuning" 文件夹复制到 Claude Code 的 skills 目录中。
- 重启 Claude Code 让新的 skill 生效。
Claude Code — 手动复制安装
- 安装前请先查看来源仓库和风险报告。
- 从源仓库下载 SKILL.md 及相关文件。
- 在 Claude Code 的 skills 目录中创建新文件夹。
- 将所有 skill 文件复制到新文件夹中。
- 重启 Claude Code 让新的 skill 生效。
Cursor — Git Clone 安装
- 安装前请先查看来源仓库和风险报告。
- 克隆仓库：git clone https://github.com/NVIDIA/skills.git
- 将 "skills/nemo-mbridge-perf-memory-tuning" 文件夹复制到 Cursor 的 skills 目录中。
- 重启 Cursor 让新的 skill 生效。
Cursor — 手动复制安装
- 安装前请先查看来源仓库和风险报告。
- 从源仓库下载 SKILL.md 及相关文件。
- 在 Cursor 的 skills 目录中创建新文件夹。
- 将所有 skill 文件复制到新文件夹中。
- 重启 Cursor 让新的 skill 生效。
GitHub Copilot — Git Clone 安装
- 安装前请先查看来源仓库和风险报告。
- 克隆仓库：git clone https://github.com/NVIDIA/skills.git
- 将 "skills/nemo-mbridge-perf-memory-tuning" 文件夹复制到 GitHub Copilot 的 skills 目录中。
- 重启 GitHub Copilot 让新的 skill 生效。
GitHub Copilot — 手动复制安装
- 安装前请先查看来源仓库和风险报告。
- 从源仓库下载 SKILL.md 及相关文件。
- 在 GitHub Copilot 的 skills 目录中创建新文件夹。
- 将所有 skill 文件复制到新文件夹中。
- 重启 GitHub Copilot 让新的 skill 生效。
Windsurf — Git Clone 安装
- 安装前请先查看来源仓库和风险报告。
- 克隆仓库：git clone https://github.com/NVIDIA/skills.git
- 将 "skills/nemo-mbridge-perf-memory-tuning" 文件夹复制到 Windsurf 的 skills 目录中。
- 重启 Windsurf 让新的 skill 生效。
Windsurf — 手动复制安装
- 安装前请先查看来源仓库和风险报告。
- 从源仓库下载 SKILL.md 及相关文件。
- 在 Windsurf 的 skills 目录中创建新文件夹。
- 将所有 skill 文件复制到新文件夹中。
- 重启 Windsurf 让新的 skill 生效。
查看 SKILL.md 原文
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:
- Set
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:Truefirst. This fixes fragmentation-induced OOM with zero performance cost. Most Slurm launch templates already include it. - 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. - Add selective activation recompute (
recompute_modules=[core_attn]) if not already enabled. See @skills/nemo-mbridge-perf-activation-recompute/SKILL.md. - Avoid increasing TP as a memory fix — doubling TP dramatically increases NVLink all-reduce volume and often kills throughput (-28% on Llama3 70B).
- Avoid increasing PP at the cost of DP — halving DP doubles gradient accumulation steps and hurts throughput (~6%).
- Consider
mlprecompute if still OOM. Saves ~3 GB but costs ~16% GPU utilization on large dense models (Llama3 70B). - 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 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:Trueis incompatible with--use-nccl-ub(NCCL user-buffer registration). See Megatron-FSDP docs.- When using CUDA graphs with
expandable_segments:True, setNCCL_GRAPH_REGISTER=0(required on pre-Blackwell GPUs, enforced by MCoreCudaGraphManager). - CPU offloading requires
pipeline_model_parallel_size = 1. - Distributed optimizer requires
use_distributed_optimizer = Truein the optimizer config. sequence_parallel_input_regatherapplies 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:Trueis 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
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)
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)
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
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
Verification
Quick check that expandable_segments:True is active:
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:
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
评论 (0)
暂无评论，成为第一个评论者吧！