[BUG] CPU latency makes MFSDP v2 slower than v1 on a DeepSeek-V3 proxy

**Describe the bug**

Megatron-FSDP v2 costs about 18% more wall-clock per training step than v1 at a
matched configuration and matched accuracy, on a DeepSeek-V3-shaped MLA/MoE proxy.
v2 uses ~33% less allocated memory, so this may be an acceptable trade -- filing it
so the cost is on record rather than as a defect claim. @NVIDIA/mcore-oncall

**Steps/Code to reproduce bug**

8xH100. The proxy below is ~0.3B params: 1 dense + 3 MoE layers, hidden 1024,
dense FFN 2816, 32 experts top-8, MoE FFN 512, one shared expert, MLA with
`q_lora_rank` 512, sigmoid router with expert bias and `seq_aux_loss` balancing,
grouped GEMM, alltoall dispatch, RMSNorm, SwiGLU, `qk_layernorm`, `seq_length` 512.
Mock data and `NullTokenizer`, so no dataset or tokenizer files are needed.

Key settings: EP=2, `global_batch_size=16`, `micro_batch_size=1` (2 microbatches),
`num_distributed_optimizer_instances=4`, `outer_dp_sharding_strategy="optim"`,
`data_parallel_sharding_strategy="optim_grads_params"`, `overlap_grad_reduce=False`,
`clip_grad=1.0`, `use_precision_aware_optimizer=True`, `train_iters=20`.

```bash
torchrun --nproc_per_node=8 pretrain_dsv3_proxy_mfsdp_v2.py \
ddp.megatron_fsdp_version=1 model.cuda_graph_impl=none
torchrun --nproc_per_node=8 pretrain_dsv3_proxy_mfsdp_v2.py \
ddp.megatron_fsdp_version=2 model.cuda_graph_impl=none
```

Compare `elapsed time per iteration (ms)` over iterations 2-20. Iteration 1 includes
startup and is not representative.

The driver is a Megatron-Bridge script, so reproducing needs Megatron-Bridge in
addition to Megatron-LM. v2 additionally needs
https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/5933, which drops stale
MFSDP v2 guard rails and installs `no_sync_func` for v2; without it the config is
rejected before training starts.

pretrain_dsv3_proxy_mfsdp_v2.py

```python
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Small DeepSeek-V3-shaped MLA/MoE proxy for sanity-testing Megatron-FSDP v2 with EP.

The architecture mirrors what the MLPerf DeepSeek-V3 671B FSDP config trains
(MLA attention, sigmoid/expert-bias router with seq_aux_loss balancing, grouped-GEMM
experts, one shared expert, flex/hybridep dispatch, a leading dense layer), scaled down so
that the whole thing fits comfortably on 8 GPUs with mock data and no tokenizer files.

Gradient clipping (#6489), the precision-aware optimizer (#6506), hybrid FSDP (#6873)
and full-iteration CUDA graphs are all turned on here, so a clean run covers the whole
MLPerf feature set except MXFP8. Everything MFSDP v2 still rejects (FP8/MXFP8,
delay_wgrad_compute, TP/PP/CP/VPP, checkpointing) is left off.

Verified together on 8xH200: loss matches an eager alltoall run to 4-5 significant
digits. Capture costs ~2 min on the first iteration, after which steps run well under
a second, so this is no longer a sub-minute script.

Three PRs are still unmerged, so the defaults here do not run against stock
checkouts:

- Megatron-Bridge `#5933 `_
(open) drops the stale v2 guard rails. On Bridge main, ``clip_grad``,
``use_precision_aware_optimizer`` and hybrid FSDP each raise in
``_validate_and_apply_megatron_fsdp_v2_configs``, so the config is rejected before
training starts.
- Megatron-LM `#7074 `_ (open) fixes
``use_precision_aware_optimizer=True`` silently disabling ``clip_grad`` under v2:
the decoupled-grad heuristic in ``clip_grad_norm`` reads ``param.decoupled_grad``,
which v2 never populates, and ``clip_grad_by_total_norm_fp32`` then skips every
parameter without erroring. Without it the two settings below are mutually
defeating and ``clip_grad`` is decorative -- the loss curve is bit-identical to a
``clip_grad=0.0`` run.
- Megatron-LM `#7075 `_ (open) adds
full-iteration CUDA graph support to MFSDP v2.

Hopper or newer is also required, because the flex dispatcher's hybridep backend runs
on DeepEP, which refuses B100.

Run it with torchrun; any field of the ConfigContainer can be overridden with
Hydra-style ``section.field=value`` arguments, for example::

torchrun --nproc_per_node=8 pretrain_dsv3_proxy_mfsdp_v2.py \
model.expert_model_parallel_size=2 train.train_iters=20

Full-iteration CUDA graphs on MoE
---------------------------------

MoE expert dispatch is normally uncapturable because its buffers are sized from a
per-step device query. The settings below exist only to remove that, and each was
checked by turning it off:

- ``moe_expert_rank_capacity_factor`` is what makes dispatch shapes static. Without
it DeepEP sizes its receive buffer per step and capture dies with
``cudaErrorStreamCaptureInvalidated``. It is a throughput knob, not a correctness
one: a runner reruns any over-budget step unpadded, so too small a value costs
reruns rather than dropped tokens.
- ``use_transformer_engine_op_fuser`` is required by the capacity factor on the
hybridep backend. ``moe_use_grouped_tensor=True`` satisfies the same check in
principle, but needs a TE whose ``GroupedLinear`` takes ``use_grouped_tensor``
(TE 2.18 does not).
- ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`` is undocumented and easy to miss: without it
the op fuser asserts "Fused GroupedMLP is not supported for this configuration"
on any SwiGLU MoE, which is every DeepSeek-V3 shape.
- ``flex``/``hybridep`` is mandatory. The ``alltoall`` dispatcher sets a CUDA sync
point unconditionally, so it can never be captured at any EP.
- ``use_te_rng_tracker=True`` and ``check_for_nan_in_loss=False`` are asserted by
full_iteration capture itself: dropout RNG has to come from TE's tracker, and the
NaN check syncs on the loss.
- ``moe_paged_stash=True`` is optional -- it only pages backward activations. Runs
with and without it produced bit-identical losses, so it is left off.
"""

from __future__ import annotations

import argparse
import os
import sys
from functools import partial

# Must be set before Transformer Engine is imported. The TE op fuser, which the
# capacity factor requires on the hybridep backend, refuses to build a fused
# GroupedMLP for SwiGLU without it -- and every DeepSeek-V3 shape is SwiGLU.
os.environ.setdefault("NVTE_CUTEDSL_FUSED_GROUPED_MLP", "1")

import torch
import torch.nn.functional as F
from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec

from megatron.bridge.models.mla_provider import MLAModelProvider
from megatron.bridge.training.config import (
CheckpointConfig,
ConfigContainer,
DistributedDataParallelConfig,
DistributedInitConfig,
LoggerConfig,
MockGPTDatasetConfig,
OptimizerConfig,
RerunStateMachineConfig,
RNGConfig,
SchedulerConfig,
TokenizerConfig,
TrainingConfig,
ValidationConfig,
runtime_config_update,
)
from megatron.bridge.training.gpt_step import forward_step
from megatron.bridge.training.pretrain import pretrain
from megatron.bridge.training.utils.omegaconf_utils import process_config_with_overrides

SEQ_LENGTH = 512
VOCAB_SIZE = 32000
NUM_LAYERS = 4
NUM_DENSE_LAYERS = 1

def dsv3_proxy_model() -> MLAModelProvider:
"""DeepSeek-V3-shaped MLA/MoE model, scaled down to a few hundred million parameters."""
return MLAModelProvider(
# get_gpt_decoder_block_spec is what the DeepSeekV3Bridge uses; the default
# GPT layer spec builds a single uniform layer and ignores moe_layer_freq.
transformer_layer_spec=partial(get_gpt_decoder_block_spec, use_transformer_engine=True),
# Shape. DeepSeek-V3 671B values are in the trailing comments.
num_layers=NUM_LAYERS, # 61
hidden_size=1024, # 7168
ffn_hidden_size=2816, # 18432 (dense layers only)
num_attention_heads=8, # 128
seq_length=SEQ_LENGTH, # 4096 in the MLPerf config
vocab_size=VOCAB_SIZE, # 129280
make_vocab_size_divisible_by=128,
# MLA. Head dims are already small in DeepSeek-V3, so they are kept as-is.
multi_latent_attention=True,
q_lora_rank=512, # 1536
kv_lora_rank=512, # 512
qk_head_dim=128, # 128
qk_pos_emb_head_dim=64, # 64
v_head_dim=128, # 128
qk_layernorm=True,
# MoE. One leading dense layer stands in for first_k_dense_replace=3.
num_moe_experts=32, # 256
moe_ffn_hidden_size=512, # 2048
moe_shared_expert_intermediate_size=512, # 2048 (1 shared expert)
moe_router_topk=8, # 8
moe_layer_freq=[0] * NUM_DENSE_LAYERS + [1] * (NUM_LAYERS - NUM_DENSE_LAYERS),
moe_grouped_gemm=True,
# flex/hybridep is what our MLPerf config runs, and it is also the only
# dispatcher that can be CUDA-graph captured: alltoall sets a CUDA sync point
# unconditionally to size its buffers. Needs Hopper or newer; DeepEP refuses
# B100. See the CUDA graph section of the module docstring.
moe_token_dispatcher_type="flex",
moe_flex_dispatcher_backend="hybridep",
# Declares a per-rank token budget so dispatch buffers are pre-sized instead of
# queried from the device each step. That static sizing is what makes capture
# possible. A runner reruns any over-budget step unpadded, so a value that is
# too small costs reruns, not dropped tokens.
moe_expert_rank_capacity_factor=1.0,
use_transformer_engine_op_fuser=True,
moe_router_load_balancing_type="seq_aux_loss",
moe_aux_loss_coeff=1e-4,
moe_router_score_function="sigmoid",
moe_router_enable_expert_bias=True,
moe_router_pre_softmax=True,
moe_router_dtype="fp32",
moe_permute_fusion=True,
# Overlapping shared-expert compute with dispatch is orthogonal to FSDP but adds
# another moving part; turn it on only once the base run is green.
moe_shared_expert_overlap=False,
moe_router_force_load_balancing=False,
# Parallelism. MFSDP v2 requires TP=PP=CP=1 and no virtual pipeline.
tensor_model_parallel_size=1,
pipeline_model_parallel_size=1,
context_parallel_size=1,
virtual_pipeline_model_parallel_size=None,
expert_model_parallel_size=2,
expert_tensor_parallel_size=1,
sequence_parallel=False,
# Precision. MFSDP v2 is BF16-only today.
bf16=True,
fp16=False,
params_dtype=torch.bfloat16,
pipeline_dtype=None,
# Rejected by MFSDP v2, spelled out so a stray default cannot turn them on.
gradient_accumulation_fusion=False,
calculate_per_token_loss=False,
# Full-iteration capture (#7075). The first iteration pays ~2 min of capture,
# then steps drop to well under a second.
cuda_graph_impl="full_iteration",
use_te_rng_tracker=True,
# Remaining DeepSeek-V3 settings.
normalization="RMSNorm",
activation_func=F.silu,
gated_linear_unit=True,
add_bias_linear=False,
# multi_latent_attention makes GPTModel skip its own rotary embedding, so
# position_embedding_type only matters here to keep it from building a learned
# absolute table; MLA picks its RoPE flavour from rope_type. DeepSeek-V3 uses
# yarn, but plain rope keeps the proxy free of long-context scaling terms.
position_embedding_type="rope",
rope_type="rope",
rotary_base=10000.0,
apply_rope_fusion=False,
share_embeddings_and_output_weights=False,
attention_dropout=0.0,
hidden_dropout=0.0,
attention_softmax_in_fp32=False,
bias_activation_fusion=True,
bias_dropout_fusion=True,
masked_softmax_fusion=True,
persist_layer_norm=True,
# Orthogonal to FSDP and version-sensitive in TE; left off so a failure here
# cannot be mistaken for an MFSDP problem.
cross_entropy_loss_fusion=False,
init_method_std=0.02,
layernorm_epsilon=1e-5,
transformer_impl="transformer_engine",
mtp_num_layers=None,
)

def build_config() -> ConfigContainer:
"""Assemble the full training configuration for the MFSDP v2 EP smoke run."""
train_iters = 20

return ConfigContainer(
model=dsv3_proxy_model(),
dist=DistributedInitConfig(use_megatron_fsdp=True),
train=TrainingConfig(
train_iters=train_iters,
# Must be divisible by data_parallel_size * micro_batch_size. On 8 GPUs with
# TP=PP=CP=1 that is 8, so this runs 2 gradient-accumulation microbatches.
global_batch_size=16,
micro_batch_size=1,
exit_signal_handler=True,
),
validation=ValidationConfig(eval_interval=train_iters + 1, eval_iters=0),
optimizer=OptimizerConfig(
optimizer="adam",
lr=1e-4,
min_lr=1e-5,
weight_decay=0.01,
adam_beta1=0.9,
adam_beta2=0.95,
adam_eps=1e-8,
bf16=True,
fp16=False,
# MFSDP v2 owns its sharded parameter/gradient storage, so the distributed
# optimizer stays off. Clipping (#6489) and the precision-aware optimizer
# (#6506) are supported now; the latter needs #7074, without which it
# silently disables clipping.
clip_grad=1.0,
use_distributed_optimizer=False,
use_precision_aware_optimizer=True,
),
scheduler=SchedulerConfig(
lr_decay_style="cosine",
lr_decay_iters=train_iters,
lr_warmup_iters=2,
lr_warmup_init=0.0,
start_weight_decay=0.01,
end_weight_decay=0.01,
weight_decay_incr_style="constant",
override_opt_param_scheduler=True,
),
ddp=DistributedDataParallelConfig(
use_megatron_fsdp=True,
megatron_fsdp_version=2,
data_parallel_sharding_strategy="optim_grads_params",
# Hybrid FSDP (#6873): shard the optimizer state over the outer DP axis, as
# our MLPerf config does. Set outer_dp_sharding_strategy=no_shard for plain
# HSDP. At 8 GPUs with EP=2 this gives the dense grid outer=4 / inner=2, the
# same ratio the MLPerf config runs at 256 GPUs. Only dense parameters go
# hybrid: v2 always puts expert parameters on a 1-D mesh over the whole
# expert-DP domain, so this knob does not change expert sharding.
outer_dp_sharding_strategy="optim",
num_distributed_optimizer_instances=4,
use_distributed_optimizer=False,
average_in_collective=False,
# Precision comes from megatron_fsdp_main_{params,grads}_dtype, which build the
# MixedPrecisionPolicy; grad_reduce_in_fp32 is inert on v2. Defaults are FP32
# main weights and parameter-dtype main grads, matching the v2 unit tests.
),
dataset=MockGPTDatasetConfig(
random_seed=1234,
seq_length=SEQ_LENGTH,
reset_attention_mask=False,
reset_position_ids=False,
eod_mask_loss=False,
num_dataset_builder_threads=1,
data_sharding=True,
dataloader_type="single",
num_workers=1,
),
# log_params_norm walks the parameters every step, which is an extra thing to get
# right against sharded storage; flip it on once the loss curve looks sane.
logger=LoggerConfig(log_interval=1, log_params_norm=False),
tokenizer=TokenizerConfig(tokenizer_type="NullTokenizer", vocab_size=VOCAB_SIZE),
# MFSDP v2 has no checkpoint support yet, so save/load stay unset.
checkpoint=CheckpointConfig(ckpt_format="fsdp_dtensor", async_save=False),
rng=RNGConfig(seed=1234),
# full_iteration capture asserts this is off: the NaN check syncs on the loss.
rerun_state_machine=RerunStateMachineConfig(check_for_nan_in_loss=False),
)

def dry_run(cfg: ConfigContainer, world_size: int) -> None:
"""Run the config validation and print the resolved config, without touching a GPU.

This is the cheapest way to see the MFSDP v2 guard rails fire, since
``runtime_config_update`` is where ``_validate_and_apply_megatron_fsdp_v2_configs``
runs.
"""
environment = {"WORLD_SIZE": str(world_size), "RANK": "0"}
previous = {name: os.environ.get(name) for name in environment}
try:
os.environ.update(environment)
runtime_config_update(cfg)
finally:
for name, value in previous.items():
if value is None:
os.environ.pop(name, None)
else:
os.environ[name] = value
cfg.print_yaml()

def main(argv: list[str] | None = None) -> None:
"""Apply CLI overrides to the smoke configuration and run training."""
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--dry-run",
dest="dryrun",
action="store_true",
help="Validate and print the resolved configuration, then exit without training.",
)
parser.add_argument(
"--world-size",
type=int,
default=8,
help="World size assumed by --dry-run. Ignored when actually training.",
)
args, overrides = parser.parse_known_args(argv)
for override in overrides:
if "=" not in override:
parser.error(f"Unrecognized argument {override!r}; overrides must be section.field=value.")

cfg = build_config()
cfg = process_config_with_overrides(cfg, cli_overrides=list(overrides) or None)

if args.dryrun:
dry_run(cfg, args.world_size)
return

pretrain(cfg, forward_step)

if __name__ == "__main__":
main(sys.argv[1:])
```

**Expected behavior**

v2 step time within noise of v1, or a documented reason for the gap.

**Additional context**

Iteration time, median over iterations 2-20:

| | v1 | v2 | delta |
| --- | --- | --- | --- |
| median | 207.1 ms | 245.1 ms | +18.4% |
| mean | 211.7 ms | 246.8 ms | +16.6% |
| min-max | 204.0-230.8 | 242.2-259.1 | ranges do not overlap |

Memory, rank 0, steady state:

| | v1 | v2 | delta |
| --- | --- | --- | --- |
| allocated | 1.2539 GB | 0.8372 GB | -33.2% |
| max-allocated | 1.6180 GB | 1.3396 GB | -17.2% |
| reserved | 1.7637 GB | 1.8476 GB | +4.8% |

Accuracy is matched, so neither version is trading one for the other. Against a
no-FSDP reference (plain DDP plus distributed optimizer, same model and the same
two microbatches), the maximum relative lm-loss deviation over 20 iterations is
4.105e-03 for v1 and 3.228e-04 for v2.

The numbers above are eager. Enabling full-iteration CUDA graphs makes step time
much **slower** for both versions -- 875.3 ms for v1 and 926.6 ms for v2, roughly
4x the eager step time -- while leaving losses and memory unchanged, with no
per-step recapture in the logs. Still investigating; not the subject of this issue.

Environment: 8xH100, `nvcr.io/nvidian/nemo:nightly`. Megatron-LM commit 2371e9022,
Megatron-Bridge commit .

Caveats: single run per configuration; rank 0 memory only.

Start with the embedded pretrain_dsv3_proxy_mfsdp_v2.py driver and run both torchrun commands on 8xH100, using the listed Megatron-Bridge and Megatron-LM pull requests where required. Compare elapsed time per iteration for iterations 2-20, then investigate the reported CPU latency difference between MFSDP v1 and v2. Done means a reproducible comparison and a documented explanation or improvement for the observed gap.

Written by the indexing model from the issue text.