Bug: deepseekv3_pretrain.yaml fails on first forward with DeepEP dispatch timeout

Describe the bug

DeepSeek-V3 pretraining fails on the first forward pass when using the examples/llm_pretrain/deepseekv3_pretrain.yaml recipe with expert parallelism and DeepEP dispatch. The run times out in DeepEP during token dispatch before any optimizer step, checkpointing, or validation occurs.

Steps/Code to reproduce bug

Use examples/llm_pretrain/deepseekv3_pretrain.yaml, unchanged except for real Megatron dataset paths.

Launch on 16 nodes x 8 GPUs, for a total of 128 GPUs, matching the configured world size.

Relevant config:

model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_config
  config:
    pretrained_model_name_or_path: deepseek-ai/DeepSeek-V3

distributed:
  _target_: nemo_automodel.components.distributed.fsdp2.FSDP2Manager
  tp_size: 1
  cp_size: 1
  pp_size: 8
  dp_size: 16
  ep_size: 16

optimizer:
  _target_: torch.optim.AdamW

The failure occurs at Training: 0%, during the first forward pass.

Error:

RuntimeError: DeepEP error: timeout (dispatch CPU)

Relevant stack locations:

nemo_automodel/components/moe/megatron/token_dispatcher.py
nemo_automodel/components/moe/megatron/fused_a2a.py
deep_ep/buffer.py: internode_dispatch

Expected behavior

The DeepSeek-V3 pretraining run should successfully complete the first forward pass and begin training without a DeepEP dispatch timeout.

Additional context

The failure happens before any optimizer step, checkpoint, or validation.

This was reproduced with both stock torch.optim.AdamW using bf16 and a TE FusedAdam variant. Both configurations fail with the same DeepEP timeout, so the issue does not appear to be optimizer- or dtype-related.

Full Slurm log attached:

te_check_deepseekv3_torch_adamw_bf16-11910803.txt

Start with examples/llm_pretrain/deepseekv3_pretrain.yaml and reproduce the first forward pass on the stated 128-GPU configuration. Read nemo_automodel/components/moe/megatron/token_dispatcher.py, nemo_automodel/components/moe/megatron/fused_a2a.py, and deep_ep/buffer.py around internode_dispatch. Done means the run completes its first forward pass without the DeepEP dispatch timeout.