NVIDIA-NeMo / NVIDIA-NeMo/RL

Split path: support MTP + SEQUENCE_LEVEL loss with mtp_detach_heads=False

Open
#3,889 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 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

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.