NVIDIA-NeMo / NVIDIA-NeMo/RL

Qwen3.5 MoE Automodel path does not use FP32 master weights by default, causing slower convergence

Open
#2,322 0 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

# Summary

On the Automodel (FSDP2) path, Qwen3.5 MoE (and other custom MoE models) end up with the optimizer operating on BF16 parameters rather than FP32 master weights. With `torch.optim.AdamW`, updates of magnitude `lr · m / (√v + ε)` at typical RLHF learning rates are well below the BF16 quantization step and get rounded to zero on most parameters. Convergence lags the MCore (distributed optimizer) path, which keeps FP32 master weights externally.

The root-cause fix belongs in `nemo-automodel` and is tracked separately — see https://github.com/NVIDIA-NeMo/Automodel/pull/1896.

This issue tracks the NeMo-RL side:

1. Document the symptom so users picking the automodel path for MoE models are aware.
2. Ship the FusedAdam workaround (#2320) for the affected recipes, which bypasses the missing FP32 master weights by having the optimizer hold its own FP32 master copy internally.

# Symptoms

Same workload (Qwen3.5-35B-A3B-Base, DAPO-style GRPO, BF16), three setups:

- Yellow: MCore path
- Red: Automodel + `torch.optim.AdamW` (default)
- Grey: Automodel + TE FusedAdam (`master_weights=True`)

Image

# Affected recipes (NeMo-RL side)

Automodel-path GRPO recipes for Qwen3.5 MoE:

- `examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-2n8g-automodel-ep16.yaml`
- `examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.yaml`
- `examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16.yaml`

(`-megatron` recipes go through the MCore distributed optimizer and are not affected.)

# Workaround (#2320)

Switch the affected recipes to `transformer_engine.pytorch.optimizers.fused_adam.FusedAdam` with:

```yaml
policy:
optimizer:
name: "transformer_engine.pytorch.optimizers.fused_adam.FusedAdam"
kwargs:
lr: ...
weight_decay: ...
betas: [0.9, 0.999]
eps: 1.0e-08
master_weights: true
store_param_remainders: true
exp_avg_dtype: "torch.bfloat16"
exp_avg_sq_dtype: "torch.bfloat16"
```

FusedAdam holds an internal FP32 master copy of the parameters, so the BF16 parameter storage on the Automodel MoE path no longer matters for optimizer precision.

# Resolution criteria

- [x] Short-term: FusedAdam enabled on the affected recipes (#2320).
- [ ] Long-term: root-cause fix in `nemo-automodel` lands and the automodel MoE path produces FP32 master weights with `torch.optim.AdamW` https://github.com/NVIDIA-NeMo/Automodel/pull/1896. Once that ships, the FusedAdam override in these recipes can be revisited (switch back to AdamW, or keep FusedAdam on perf grounds).

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.