Split path: MoE auxiliary loss inherits the main loss denominator instead of 1/global_valid_toks
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Summary
On the SingleController split training path the MoE auxiliary loss is normalized by the **main loss's** denominator rather than by the token count, so it diverges from the synchronous `train()` path whenever the loss is `SEQUENCE_LEVEL`.
`_finish_train_step_body` applies one uniform [`self.model.scale_gradients(inv_n)`](https://github.com/NVIDIA-NeMo/RL/blob/b0dd48137db9913db3196a13f69bbd1ea10572cb/nemo_rl/models/policy/workers/megatron_policy_worker.py#L1640), where [`inv_n`](https://github.com/NVIDIA-NeMo/RL/blob/b0dd48137db9913db3196a13f69bbd1ea10572cb/nemo_rl/models/policy/workers/megatron_policy_worker.py#L1629-L1634) is `1/global_valid_toks` for `TOKEN_LEVEL` but `1/global_valid_seqs` for `SEQUENCE_LEVEL`. Every auxiliary branch inherits it.
[`_set_moe_grad_scale_func`](https://github.com/NVIDIA-NeMo/RL/blob/b0dd48137db9913db3196a13f69bbd1ea10572cb/nemo_rl/models/policy/workers/megatron_policy_worker.py#L1157) is only ever called inside the synchronous `train()` — [:910](https://github.com/NVIDIA-NeMo/RL/blob/b0dd48137db9913db3196a13f69bbd1ea10572cb/nemo_rl/models/policy/workers/megatron_policy_worker.py#L910) and [:953](https://github.com/NVIDIA-NeMo/RL/blob/b0dd48137db9913db3196a13f69bbd1ea10572cb/nemo_rl/models/policy/workers/megatron_policy_worker.py#L953). The split path never sets the hook, so Megatron falls back to `grad_scale_func` via [`_get_moe_loss_scale`](https://github.com/NVIDIA/Megatron-LM/blob/731b791469004f8fdcb896e65d610d6b1b0ebb32/megatron/core/pipeline_parallel/schedules.py#L263-L268) → [`_compute_loss_scale`](https://github.com/NVIDIA/Megatron-LM/blob/731b791469004f8fdcb896e65d610d6b1b0ebb32/megatron/core/pipeline_parallel/schedules.py#L254-L260). The fallback itself is upstream-sanctioned — [`ModelParallelConfig.moe_grad_scale_func`](https://github.com/NVIDIA/Megatron-LM/blob/731b791469004f8fdcb896e65d610d6b1b0ebb32/megatron/core/model_parallel_config.py#L200) documents `None` as "falls back to `grad_scale_func`" — so the issue is the resulting denominator, not the mechanism.
Legacy pins the MoE aux gradient to `1/global_valid_toks` via [`_compute_moe_grad_scale`](https://github.com/NVIDIA-NeMo/RL/blob/b0dd48137db9913db3196a13f69bbd1ea10572cb/nemo_rl/models/policy/workers/megatron_policy_worker.py#L1146-L1155), and the derivation in the sync call-site comment ([:902-909](https://github.com/NVIDIA-NeMo/RL/blob/b0dd48137db9913db3196a13f69bbd1ea10572cb/nemo_rl/models/policy/workers/megatron_policy_worker.py#L902-L909)) spells out the per-token reasoning.
**Effect:** under a `SEQUENCE_LEVEL` loss the MoE aux term is weighted by mean-tokens-per-sequence relative to legacy. Since response length drifts during RL training, the effective `moe_aux_loss_coeff` becomes length-dependent rather than constant, so a value tuned at one response length does not transfer.
`TOKEN_LEVEL` is exactly equivalent between the two paths — both come out at `1/global_valid_toks`. Nothing to do there.
## Why it is harder than the MTP case
There is no separable-parameter escape. The MoE aux-loss autograd trick routes the gradient through the router and everything upstream of it, so no parameter subset carries only the aux term — unlike MTP with `mtp_detach_heads=True`, where the params are tagged `grad_norm_group='mtp'` and can be rescaled independently at finish.
Pre-compensating at backward time would require `n_true / global_valid_toks`, and neither global exists until the finish-time all-reduce: chunks are streamed and processed as they arrive, so no per-chunk constant can carry it.
## Possible directions
- Make the step's global token/sequence counts available before the backward passes (changes the streaming contract).
- Give auxiliary branches their own deferred normalizer rather than sharing the main loss's `inv_n`.
- Refuse `SEQUENCE_LEVEL` + MoE on the split path until one of the above lands.
## Notes
`SEQUENCE_LEVEL` is not only an opt-in normalization knob: GSPO requires it by construction — [`ClippedPGLossFn.__init__`](https://github.com/NVIDIA-NeMo/RL/blob/b0dd48137db9913db3196a13f69bbd1ea10572cb/nemo_rl/algorithms/loss/loss_functions.py#L262-L268) asserts `loss_type == SEQUENCE_LEVEL` when `sequence_level_importance_ratios=True`. `PreferenceLossFn` / `DPOLossFn` are hard-coded `SEQUENCE_LEVEL`.
Surfaced during review of #3881, which ports the MTP training lifecycle to the split path. That PR handles the MTP half; this issue tracks MoE, which is pre-existing and independent.
Related: #3889 tracks the MTP half (`mtp_detach_heads=False`), which is separable-in-principle for the detached case but not for the attached one.
Contributor guide
Assessment
This issue has not been assessed yet.