Nemo Mbridge Perf Megatron Fsdp

nvidia/skills/skills/nemo-mbridge-perf-megatron-fsdp

by nvidiacf5224d14250Apache-2.03.5K starsListed Oct 8, 2026Updated Oct 8, 2026Repository updated today

Operational guide for enabling Megatron FSDP in Megatron-Bridge, including config knobs, code anchors, pitfalls, and verification.

Instructions onlyDevOps & Cloud

AI-generated overview

Operational guide for enabling Megatron FSDP in Megatron-Bridge, covering config knobs, code anchors, pitfalls, and verification.

What it does: This skill provides an operational guide for turning on Megatron FSDP within Megatron-Bridge. It lists the required configuration overrides, points to relevant code anchors in the Bridge config and runtime wrapper, describes common pitfalls, and gives a verification procedure using an existing 2-GPU smoke test. It produces guidance and configuration snippets rather than executable artifacts.
When to use it: Use it when switching from DDP to FSDP-based data parallelism, or when tracing an out-of-memory error or performance regression to an FSDP configuration change. It is also relevant when working with settings such as use_megatron_fsdp, data_parallel_sharding_strategy, or sharded data parallel.
Requirements: Requires access to the Megatron-Bridge codebase and its documentation. Verification uses pytest, torch.distributed, and uv, and needs at least two GPUs. No scripts are shipped; the skill is instructions only.

Megatron 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:

python

cfg.dist.use_megatron_fsdp = Truecfg.ddp.use_megatron_fsdp = Truecfg.ddp.data_parallel_sharding_strategy = "optim_grads_params"cfg.ddp.average_in_collective = Falsecfg.checkpoint.ckpt_format = "fsdp_dtensor"

cfg.dist.use_megatron_fsdp = Truecfg.ddp.use_megatron_fsdp = Truecfg.ddp.data_parallel_sharding_strategy = "optim_grads_params"cfg.ddp.average_in_collective = Falsecfg.checkpoint.ckpt_format = "fsdp_dtensor"

Example recipe fixup:

python

cfg = llama3_8b_pretrain_config()cfg.dist.use_megatron_fsdp = Truecfg.ddp.use_megatron_fsdp = Truecfg.ddp.data_parallel_sharding_strategy = "optim_grads_params"cfg.ddp.average_in_collective = Falsecfg.checkpoint.ckpt_format = "fsdp_dtensor"cfg.checkpoint.save = "/tmp/fsdp_ckpts"cfg.checkpoint.load = None

cfg = llama3_8b_pretrain_config()cfg.dist.use_megatron_fsdp = Truecfg.ddp.use_megatron_fsdp = Truecfg.ddp.data_parallel_sharding_strategy = "optim_grads_params"cfg.ddp.average_in_collective = Falsecfg.checkpoint.ckpt_format = "fsdp_dtensor"cfg.checkpoint.save = "/tmp/fsdp_ckpts"cfg.checkpoint.load = None

Performance harness note:

Code Anchors

Bridge config definition:

148:154:src/megatron/bridge/training/config.py

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 otherpotential issues."""

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 otherpotential issues."""

Bridge validation:

1533:1578:src/megatron/bridge/training/config.py

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")

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:

217:243:src/megatron/bridge/models/common/unimodal.py

if use_megatron_fsdp:    DP = FullyShardedDataParallelelif use_torch_fsdp2:    DP = TorchFullyShardedDataParallelelse:    DP = DistributedDataParallel...DP(    config=get_model_config(model_chunk),    ddp_config=ddp_config,    module=model_chunk,    ...    pg_collection=pg_collection,)

if use_megatron_fsdp:    DP = FullyShardedDataParallelelif use_torch_fsdp2:    DP = TorchFullyShardedDataParallelelse:    DP = DistributedDataParallel...DP(    config=get_model_config(model_chunk),    ddp_config=ddp_config,    module=model_chunk,    ...    pg_collection=pg_collection,)

Perf harness overrides:

74:98:scripts/performance/utils/overrides.py

recipe.ddp.use_megatron_fsdp = Truerecipe.ddp.data_parallel_sharding_strategy = "optim_grads_params"recipe.ddp.keep_fp8_transpose_cache = Falserecipe.ddp.average_in_collective = False...recipe.checkpoint.load = None

recipe.ddp.use_megatron_fsdp = Truerecipe.ddp.data_parallel_sharding_strategy = "optim_grads_params"recipe.ddp.keep_fp8_transpose_cache = Falserecipe.ddp.average_in_collective = False...recipe.checkpoint.load = None

Pitfalls

Public recipes often expose use_megatron_fsdp but still default to ckpt_format="torch_dist". If save/load is enabled, switch to fsdp_dtensor.
use_torch_fsdp2 exists, but on the validated branch Bridge still fails before training because _ddp_wrap passes pg_collection.
CPU offloading is only valid when pipeline_model_parallel_size == 1 and activation recomputation is disabled.
Upstream warns that FSDP and TP/CP can want different CUDA_DEVICE_MAX_CONNECTIONS settings on Hopper and earlier.
Megatron FSDP and FSDP2 are mutually exclusive.

Verification

Use the existing 2-GPU functional smoke test:

bash

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

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

Source and attribution

Source:nvidia/skillsinskills/nemo-mbridge-perf-megatron-fsdpat commitcf5224d

License: Apache-2.0

Content belongs to its original authors. SourceWeft indexes it from a public repository.

Report or request removal

Nemo Mbridge Perf Megatron Fsdp

Megatron FSDP Skill

Enablement

Code Anchors

Pitfalls

Verification

Source and attribution

More from nvidia/skills

Tilegym Cutile Python

Tilegym Monkey Patch Kernels To Transformers

Tilegym Converting Cutile To Julia

Tilegym Adding Cutile Kernel

Physicsnemo Discover

Nemoclaw User Guide

More in DevOps & Cloud

Playwright Devops

M5 Onboard

Yeet

Runbook

Incident Response

Deploy Checklist