[megatron] feat: support Megatron-FSDP mode for Megatron backend

[conver334]

What does this PR do?

Support using Megatron-FSDP for SFT and RL.

Add Megatron-FSDP as a new training backend option for the Megatron engine. This is implementation of #5244 .

Key changes:

• Add example scripts for GRPO examples/grpo_trainer/run_qwen2-7b_math_megatron_fsdp.sh and SFT examples/sft/gsm8k/run_qwen_megatron_fsdp.sh. Run these following the user guide docs/examples/megatron_fsdp_example.rst
• Add use_megatron_fsdp config flag to McoreEngineConfig, enabling using Megatron-FSDP via a single config option.
• When enabled, automatically configure the required DDP settings (distributed optimizer, Zero-3 sharding strategy, grad reduce overlap) with sensible defaults, while still allowing fine-grained override via override_ddp_config.

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: [megatron, model] feat: qwen3.5 example  #5381
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • {modules} include fsdp, megatron, veomni, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data, cfg, reward, fully_async, one_step_off
  • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
  • {type} is in feat, fix, refactor, chore, test
  • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
  • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

Loss is normal in SFT. Tested on 8*H100 with Moonlight-16B-A3B-Instruct, GSM8K SFT dataset

Loss is normal in SFT. Tested on 8*H100 with Qwen2.5-Math-7B, GSM8K SFT dataset

Reward is normal in GRPO. Tested on 8*H100 with Qwen2.5-Math-7B, GSM8K

MFU

API and Usage Example

Enable Megatron-FSDP by setting three config flags:

    actor_rollout_ref.actor.megatron.use_mbridge=True \
    actor_rollout_ref.actor.megatron.vanilla_mbridge=False \
    actor_rollout_ref.actor.megatron.use_megatron_fsdp=True \

The FSDP-specific DDP settings (sharding strategy, overlap, etc.) are auto-configured with defaults. Advanced users can override them:

    actor_rollout_ref.actor.megatron.override_ddp_config.data_parallel_sharding_strategy=optim_grads \
    actor_rollout_ref.actor.megatron.override_ddp_config.overlap_grad_reduce=False \

Design & Code Changes

Megatron-FSDP use the same training loop as Megatron.
Conversion between HuggingFace format and Megatron-FSDP DTensor is implemented via NVIDIA-NeMo/Megatron-Bridge#1910.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)
• If your PR is related to the recipe submodule, please also update the reference to the submodule commit via git submodule update --remote or cd recipe && git pull origin main.

[gemini-code-assist]

Code Review

This pull request introduces support for Megatron-FSDP as a new training backend, including configuration flags, automatic DDP settings, and state lifecycle management. The overall implementation is sound, but I've identified a critical issue in the FSDP parameter synchronization logic that could lead to inconsistent model states during inference. Additionally, the new example script has a hardcoded model path, which impacts its portability. I've provided suggestions to fix these issues.

[yxs]

I tested both SFT and GRPO (without ALL_OFFLOAD) on 8×H100 with Qwen2.5-0.5B, Megatron-LM main, and Megatron-Bridge PR #1910:

• SFT: loss decreasing normally (0.695 → 0.543 over 21 steps)
• GRPO: 447 steps completed, reward improving (0.006 → 0.72), DTensor weight export works correctly

@conver334 All looks good to me now.