Megatron FSDP Configuration and Performance Guide
Enable, troubleshoot, and verify Megatron FSDP in Megatron-Bridge.
The material only guides configuration and distributed-training execution; there is no evidence of malware, credential handling, covert exfiltration, or destructive defaults. Scope and checkpoint behavior are reasonably transparent and a verification command is provided. Points are deducted because user confirmation, permission isolation, sensitive-data handling, rollback/recovery, and safeguards for substantial compute and checkpoint side effects are not specified.
Configuration knobs, incompatibilities, code anchors, pitfalls, and the 2-GPU smoke-test path are broadly consistent, with some limitations documented. Points are deducted for incomplete dependency versions, environment prerequisites, failure diagnosis, and abnormal-input handling. Static review cannot establish reproducibility, so the score remains at or below 10.
The audience, FSDP enablement use case, OOM/regression diagnosis scenario, and trigger terms are clear, and limitations involving FSDP2, TP/PP/CP, and CPU offload are stated. Points are deducted because non-fit boundaries, input/output contracts, and environment matrices are incomplete, with no Chinese-language support guidance. The core function does not depend entirely on overseas services, so no additional mainland-network deduction is applied.
The SKILL.md is readable and includes enablement examples, code anchors, pitfalls, and verification. License, NVIDIA ownership, validation/version information, and follow-up items are available in companion files. Points are deducted for missing recommended Instructions and Examples sections, absent installation/dependency notes, missing author and tags metadata, incomplete changelog information, and no clearly defined skill-level maintenance/update path.
The examples directly provide the essential configuration, checkpoint format, performance entry point, and success criteria, covering the core evaluation task. The supplied evaluation report also indicates useful positive-task correctness and effectiveness. Points are deducted because only one task was evaluated, performance benefits are marked doc_only, and compatibility across parallelism combinations still requires user-side validation; therefore the static score does not exceed 7.
Concrete source paths, configuration anchors, a functional-test path, a card.yaml evidence list, and an evaluation report provide some auditability. Points are deducted because the supplied material lacks independently reviewable test logs or complete CI coverage; the card's live-validation statements are report-level claims, and performance claims are explicitly documentation-only. The score therefore does not exceed 5.
- When checkpoint saving or loading is enabled, confirm fsdp_dtensor and validate the target branch's configuration constraints first.
- Do not enable Megatron FSDP and Torch FSDP2 together; use_tp_pp_dp_mapping and several TP/PP/CP/EP combinations require separate validation.
- The smoke-test command depends on two CUDA GPUs, uv, a PyTorch distributed environment, and project test dependencies; this review did not execute it.
- No independently reproducible performance data is supplied, so actual speedup or memory savings should not be inferred from this skill alone.
What does this skill do, and when should you use it?
This skill is for developers using Megatron FSDP data parallelism in Megatron-Bridge. It provides the required configuration overrides, code anchors, performance-harness entry point, known pitfalls, and a two-GPU functional verification command. It focuses on checkpoint format requirements, mutual exclusion with FSDP2, and configuration changes that may cause OOMs or regressions. It is most useful when moving from DDP to FSDP or investigating an FSDP-related failure.
It directs users to set use_megatron_fsdp, data_parallel_sharding_strategy, average_in_collective, and fsdp_dtensor-related settings; supplies a recipe fixup and performance-harness command; identifies the source locations for configuration definitions, validation logic, runtime data-parallel wrapper selection, and performance overrides; documents constraints involving CPU offloading, TP/CP, and CUDA_DEVICE_MAX_CONNECTIONS; and provides a two-GPU pytest smoke test with success criteria of one passing test, finite final-iteration loss, and no checkpoint-format assertion.
- An engineer needs to switch Megatron-Bridge training from DDP to Megatron FSDP.
- A developer is investigating an OOM, regression, or checkpoint failure after enabling FSDP.
- A training maintainer needs to confirm that fsdp_dtensor checkpoint settings are correct.
- A code reviewer needs anchors for Megatron-Bridge FSDP configuration, validation, or runtime wrapper selection.
- An engineer needs to run the basic FSDP training smoke test on two GPUs.
What are this skill's strengths and limitations?
- Provides copyable minimal configuration and recipe examples.
- Covers configuration, runtime selection, performance overrides, pitfalls, and verification.
- Clearly states that Megatron FSDP requires the fsdp_dtensor checkpoint format.
- Includes an existing two-GPU functional smoke-test command.
- It focuses on Megatron-Bridge Megatron FSDP rather than serving as a general FSDP tutorial.
- It does not provide complete setup instructions for Megatron-Bridge, CUDA, or the GPU environment.
- It does not specify validated GPU models, software versions, or a full platform matrix.
- FSDP2 may fail before training on the referenced branch because of the pg_collection argument; the skill does not implement a fix.
How do you install this skill?
Install the specific skill through the skills CLI used by the NVIDIA/skills repository:
npx skills add nvidia/skills --skill nemo-mbridge-perf-megatron-fsdp --yes
The README does not document the exact installation directory; the CLI handles the destination. The skill becomes available after the agent reloads its skills.
How do you use this skill?
Ask the agent for a Megatron FSDP task, for example: “Enable Megatron FSDP in Megatron-Bridge and check the checkpoint format and verification test.” The documented trigger terms include use_megatron_fsdp, data_parallel_sharding_strategy, sharded data parallel, and Megatron FSDP. The guide recommends setting cfg.dist.use_megatron_fsdp and cfg.ddp.use_megatron_fsdp to true, using optim_grads_params for data_parallel_sharding_strategy, setting average_in_collective to false, and selecting fsdp_dtensor for checkpoint.ckpt_format. The verification command is:
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
How does this skill compare with similar options?
The guide explicitly distinguishes Megatron FSDP from DDP and torch FSDP2: runtime selection chooses different data-parallel wrappers for them. Megatron FSDP and FSDP2 are mutually exclusive, and FSDP2 is currently not working with pipeline parallelism. The source provides no benchmark proving that one option is universally faster than the others.