NVIDIA-NeMo / NVIDIA-NeMo/RL

Split path: MoE auxiliary loss inherits the main loss denominator instead of 1/global_valid_toks

Open
#3,888 0 comments 0 reactions 1 assignee Claimed by @terrykong View on GitHub
accuracy bug
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

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.