Split path: support MTP + SEQUENCE_LEVEL loss with mtp_detach_heads=False
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Summary
On the SingleController split training path, the MTP auxiliary gradient is normalized by the **main loss's** denominator rather than by the token count. With `mtp_detach_heads=True` this is fixable and is being handled in #3881; with `mtp_detach_heads=False` (the default) it is not fixable on this path, so #3881 refuses the combination. This issue tracks making it actually work.
`_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`.
The MTP branch enters as a raw token **sum** — [`MTPLossAutoScaler.backward`](https://github.com/NVIDIA/Megatron-LM/blob/731b791469004f8fdcb896e65d610d6b1b0ebb32/megatron/core/transformer/multi_token_prediction.py#L734-L748) returns `torch.ones_like(mtp_loss) * scale`, seeding every element of the `[b, s]` tensor — and [`process_mtp_loss`](https://github.com/NVIDIA/Megatron-LM/blob/731b791469004f8fdcb896e65d610d6b1b0ebb32/megatron/core/transformer/multi_token_prediction.py#L875-L888) notes that with `calculate_per_token_loss=True` the expected divide is by `total_num_tokens`. So the correct denominator is `1/global_valid_toks`, which is what the synchronous `train()` applies unconditionally at [:913-915](https://github.com/NVIDIA-NeMo/RL/blob/b0dd48137db9913db3196a13f69bbd1ea10572cb/nemo_rl/models/policy/workers/megatron_policy_worker.py#L913-L915).
**Effect:** under a `SEQUENCE_LEVEL` loss the MTP term is weighted by mean-tokens-per-sequence relative to legacy. Since response length drifts during RL training, the effective `mtp_loss_scaling_factor` 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 `mtp_detach_heads=False` cannot be fixed on this path today
With `mtp_detach_heads=True`, Megatron tags every MTP parameter [`grad_norm_group='mtp'`](https://github.com/NVIDIA/Megatron-LM/blob/731b791469004f8fdcb896e65d610d6b1b0ebb32/megatron/core/transformer/multi_token_prediction.py#L1727-L1730) — note the tagging happens **only** under that flag — and the detaches at [:819-823](https://github.com/NVIDIA/Megatron-LM/blob/731b791469004f8fdcb896e65d610d6b1b0ebb32/megatron/core/transformer/multi_token_prediction.py#L819-L823), [:1144-1145](https://github.com/NVIDIA/Megatron-LM/blob/731b791469004f8fdcb896e65d610d6b1b0ebb32/megatron/core/transformer/multi_token_prediction.py#L1144-L1145) and [:1855](https://github.com/NVIDIA/Megatron-LM/blob/731b791469004f8fdcb896e65d610d6b1b0ebb32/megatron/core/transformer/multi_token_prediction.py#L1855) mean the MTP loss reaches no other parameters. Those gradients can therefore be rescaled independently at finish. #3881 does this.
With `mtp_detach_heads=False` the MTP gradient is summed into shared backbone parameters and cannot be isolated even in principle. Pre-compensating at backward time would require the factor `n_true / global_valid_toks`, and neither global exists until the finish-time all-reduce — chunks are streamed through `train_microbatches_from_meta` and processed as they arrive, so no per-chunk constant can carry it.
This is also the higher-risk half: without the `'mtp'` tag there is no separate grad-norm group, so MTP parameters are clipped **jointly** with the policy. A mis-scaled MTP term distorts the policy update rather than being absorbed by its own clip.
## 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`.
- Keep the refusal and document `mtp_detach_heads=True` as the supported configuration for sequence-level losses.
## 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`.
The MoE auxiliary branch has the same defect one branch over and is tracked separately in #3888 — it is strictly harder, since the aux gradient routes through the router and everything upstream, so no parameter subset carries only the aux term.
Surfaced during review of #3881.
Contributor guide
Assessment
This issue has not been assessed yet.