Megatron FSDP Skill

nemo-mbridge-perf-megatron-fsdp

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

nvidia/skills·1.8Kinstalls3.5Kstars

SKILL.md

Full skill instructions

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:

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

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

Install skill

Copy and run in your terminal

npx -y skills add https://github.com/nvidia/skills --skill nemo-mbridge-perf-megatron-fsdp

View on GitHub

nemo-mbridge-perf-megatron-fsdp

rag-blueprint
NVIDIA RAG Blueprint — deploy, configure, troubleshoot, and manage. Handles any RAG action: deploy, install, start, enable, disable, toggle, change, configure, troubleshoot, debug, fix, shutdown, stop, or tear down any RAG feature or service (Agentic RAG, VLM, guardrails, query rewriting, models, search, ingestion, observability, summarization, reasoning, and more).
3.1K installs
accelerated-computing-cudf
Official NVIDIA-authored guidance for NVIDIA cuDF GPU DataFrames, pandas acceleration, dask-cuDF, ETL, joins, groupby, CSV/Parquet I/O, nullable semantics, and multi-GPU DataFrame workloads.
3K installs
cuopt-developer
Modify, build, test, debug, and contribute to NVIDIA cuOpt (C++/CUDA, Python, server, CI). Use for solver internals, PRs, DCO, and code conventions.
2.7K installs
cuopt-install
Install cuOpt for Python, C, or server via pip, conda, or Docker; verify the install. For building cuOpt from source, see cuopt-developer.
2.5K installs

More skills from nvidia

nvidia/skills

cudaq-guide

CUDA-Q onboarding guide for installation, test programs, GPU simulation, QPU hardware, and quantum applications.

3.5K 2.5K

View

nvidia/skills

data-designer

Use when the user wants to create a dataset, generate synthetic data, or build a data generation pipeline.

3.5K 2.4K

View

nvidia/skills

NVIDIA DeepStream SDK development with Python pyservicemaker API. Use when building video analytics pipelines, GStreamer-based video processing, TensorRT inference integration, object detection/tracking, or Kafka/message broker integration.

3.5K 2.4K

View

nvidia/skills

rag-eval

Filesystem RAG benchmarks: corpus/, train.json, evaluate_rag.py (RAGAS quality). Not for prod monitoring, latency/throughput benchmarking (use rag-perf), or evals outside this repo layout.

3.5K 2.3K

View

nvidia/skills

nemo-retriever

Use when the user wants to search, query, extract, transcribe, describe, quote, filter, or aggregate across documents — PDFs, scanned forms / images (`.jpg` `.png` `.tiff`), Office (`.docx` `.pptx`), text (`.html` `.txt`), audio (`.mp3` `.wav` `.m4a`), or video (`.mp4` `.mov`). Prefer this over native Read / Grep for multi-file or non-PDF corpora. Not for: editing files, web browsing, single-file plain-text lookups, fine-tuning.

3.5K 2.3K

View

nvidia/skills

omniverse-cad-to-simready

Coordinate the end-to-end CAD/source-asset to SimReady workflow. Use for broad requests such as CAD to SimReady, source asset to simulation-ready USD, or prop packaging that require conversion, material/physics assignment, SimReady conformance, validation, and optional package creation; deploy or verify Content Agents services first when property assignment is enabled; route single-stage work through nested references.

3.5K 2.3K

View

nvidia/skills

deepstream-import-vision-model

Use this skill to bring any vision model from HuggingFace or NVIDIA NGC into an NVIDIA DeepStream pipeline with end-to-end automation: ONNX download, SafeTensors export, TRT engine build, custom nvinfer bbox parser, multi-stream benchmark, and PDF report. Object detection models only.

3.5K 2.3K

View

nvidia/skills

omniverse-usd-performance-tuning

Top-level workflow skill for USD performance diagnosis and optimization. Handles slow loading, high memory, low FPS, and broad scene-optimization requests; delegates auth/runtime setup to Phase 0 owners.

3.5K 2.3K

View

nvidia/skills

cuopt-routing-api-python

Vehicle routing (VRP, TSP, PDP) with cuOpt — Python API only. Use when the user is building or solving routing in Python.

3.5K 2.2K

View

nvidia/skills

skill-card-generator

Use only to generate or update a governance skill card for a specified existing agent skill directory. Do not use for explaining, listing, comparing, or discussing skill capabilities.

3.5K 2.2K

View

nvidia/skills

omniverse-realtime-viewer

Use as the top-level router for Omniverse Realtime Viewer USD app requests and focused viewer reference documents.

3.5K 2.2K

View

nvidia/skills

rag-perf

Performance benchmarking for a deployed NVIDIA RAG Blueprint server: profiling pass + aiperf load test driven by a single YAML config. Not for accuracy / RAGAS scoring (use rag-eval) or for deploying / repairing services (use rag-blueprint).

3.5K 2.1K

View