Nemo Mbridge Perf Megatron Fsdp Agent Skill
Operational guide for enabling Megatron FSDP in Megatron-Bridge, including config knobs, code anchors, pitfalls, and verification.
Install
one command, takes just this skill from the repositorynpx skills add https://github.com/NVIDIA/skills --skill nemo-mbridge-perf-megatron-fsdpWhat comes with it
16 458 bytes besides the instructionThe instruction itself
5 sections, as written by the authorMegatron FSDP Skill
For stable background and recommendation level, see:
- @docs/training/megatron-fsdp.md
- @skills/nemo-mbridge-perf-megatron-fsdp/card.yaml
Enablement
Minimal Megatron FSDP override in Bridge:
cfg.dist.use_megatron_fsdp = True
cfg.ddp.use_megatron_fsdp = True
cfg.ddp.data_parallel_sharding_strategy = "optim_grads_params"
cfg.ddp.average_in_collective = False
cfg.checkpoint.ckpt_format = "fsdp_dtensor"
Example recipe fixup:
cfg = llama3_8b_pretrain_config()
cfg.dist.use_megatron_fsdp = True
cfg.ddp.use_megatron_fsdp = True
cfg.ddp.data_parallel_sharding_strategy = "optim_grads_params"
cfg.ddp.average_in_collective = False
cfg.checkpoint.ckpt_format = "fsdp_dtensor"
cfg.checkpoint.save = "/tmp/fsdp_ckpts"
cfg.checkpoint.load = None
Performance harness note:
python scripts/performance/launch.py --use_megatron_fsdp true
Code Anchors
Bridge config definition:
use_megatron_fsdp: bool = False
"""Use Megatron's Fully Sharded Data Parallel. Cannot be used together with use_torch_fsdp2."""
use_torch_fsdp2: bool = False
"""Use the torch FSDP2 implementation. FSDP2 is not currently working with Pipeline Parallel.
It is still not in a stable release stage, and may therefore contain bugs or other
potential issues."""
Bridge validation:
if self.dist.use_megatron_fsdp and self.dist.use_torch_fsdp2:
raise ValueError(...)
...
assert not self.dist.use_tp_pp_dp_mapping, "use_tp_pp_dp_mapping is not supported with Megatron FSDP"
...
assert self.checkpoint.ckpt_format == "fsdp_dtensor", (
"Megatron FSDP only supports fsdp_dtensor checkpoint format"
)
Runtime wrapper selection:
if use_megatron_fsdp:
DP = FullyShardedDataParallel
elif use_torch_fsdp2:
DP = TorchFullyShardedDataParallel
else:
DP = DistributedDataParallel
...
DP(
config=get_model_config(model_chunk),
ddp_config=ddp_config,
module=model_chunk,
...
pg_collection=pg_collection,
)
Perf harness overrides:
recipe.ddp.use_megatron_fsdp = True
recipe.ddp.data_parallel_sharding_strategy = "optim_grads_params"
recipe.ddp.keep_fp8_transpose_cache = False
recipe.ddp.average_in_collective = False
...
recipe.checkpoint.load = None
Pitfalls
- Public recipes often expose
use_megatron_fsdpbut still default tockpt_format="torch_dist". If save/load is enabled, switch tofsdp_dtensor. use_torch_fsdp2exists, but on the validated branch Bridge still fails before training because_ddp_wrappassespg_collection.- CPU offloading is only valid when
pipeline_model_parallel_size == 1and activation recomputation is disabled. - Upstream warns that FSDP and TP/CP can want different
CUDA_DEVICE_MAX_CONNECTIONSsettings on Hopper and earlier. - Megatron FSDP and FSDP2 are mutually exclusive.
Verification
Use the existing 2-GPU functional smoke test:
CUDA_VISIBLE_DEVICES=0,1 uv run python -m torch.distributed.run --nproc_per_node=2 \
-m pytest tests/functional_tests/training/test_megatron_fsdp.py::TestMegatronFSDP::test_fsdp_pretrain_basic -v -s
Success criteria:
- Pytest reports
1 passed - The log shows finite loss at the last iteration
- The run finishes without a checkpoint format assertion
Other skills for the same job
different authors, same section of the catalogue
Guide for creating high-quality MCP (Model Context Protocol) servers that enable LLMs to interact with external services through well-designed tools. Use when building MCP servers to integrate external APIs or services, whether in Python (FastMCP) or Node/TypeScript (MCP SDK).
Automatically creates user-facing changelogs from git commits by analyzing commit history, categorizing changes, and transforming technical commits into clear, customer-friendly release notes. Turns hours of manual changelog writing into minutes of automated generation.
Use when implementation is complete, all tests pass, and you need to decide how to integrate the work - guides completion of development work by presenting structured options for merge, PR, or cleanup
Guide for creating high-quality MCP (Model Context Protocol) servers that enable LLMs to interact with external services through well-designed tools. Use when building MCP servers to integrate external APIs or services, whether in Python (FastMCP) or Node/TypeScript (MCP SDK).
React Native and Expo best practices for building performant mobile apps. Use when building React Native components, optimizing list performance, implementing animations, or working with native modules. Triggers on tasks involving React Native, Expo, mobile performance, or native platform APIs.
React and Next.js performance optimization guidelines from Vercel Engineering. This skill should be used when writing, reviewing, or refactoring React/Next.js code to ensure optimal performance patterns. Triggers on tasks involving React components, Next.js pages, data fetching, bundle optimization, or performance improvements.
Next.js best practices - file conventions, RSC boundaries, data patterns, async APIs, metadata, error handling, route handlers, image/font optimization, bundling
Use when starting feature work that needs isolation from current workspace or before executing implementation plans - creates isolated git worktrees with smart directory selection and safety verification
How to use it
Take nvidia/nemo-mbridge-perf-megatron-fsdp from the repository into ~/.claude/skills for personal
use, or into .claude/skills inside a project.
The agent identifies a skill by the name field in its header. Two skills with the
same name cannot sit side by side — one of them will be ignored.