From 7c9994f28a501bd7d4d69b560b1ff92574305d74 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Mon, 16 Feb 2026 16:57:36 -0800 Subject: [PATCH 1/2] Roll input IDs for MTP labels in RL mode When labels are not provided (RL training), create shifted labels from input_ids by rolling the tensor. Also handles None tensors in roll_tensor, fixes checkpointed_forward to handle non-tensor kwargs, and clamps num_tokens to prevent division by zero. Signed-off-by: Yi-Fu Wu --- megatron/core/models/gpt/gpt_model.py | 1 + megatron/core/models/hybrid/hybrid_model.py | 3 ++ .../transformer/multi_token_prediction.py | 33 +++++++++++++++++-- 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 99853939f4c..b2350e1ecf9 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -684,6 +684,7 @@ def _postprocess( cp_group=self.pg_collection.cp, packed_seq_params=packed_seq_params, scale_logits_fn=self._scale_logits if self.config.use_mup else None, + input_ids=input_ids, ) sequence_parallel_override = False diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index 511b24673b0..850b8f45c63 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -531,6 +531,8 @@ def forward( if in_inference_mode or is_spec_decode: self._decoder_hidden_states_cache = hidden_states else: + # For RL (labels is None), process_mtp_loss derives labels from + # input_ids to match the SFT label format. hidden_states = process_mtp_loss( hidden_states=hidden_states, labels=labels, @@ -544,6 +546,7 @@ def forward( cp_group=self.pg_collection.cp, packed_seq_params=packed_seq_params, scale_logits_fn=self._scale_logits if self.config.use_mup else None, + input_ids=input_ids, ) sequence_parallel_override = False if ( diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index d9d22587a07..348f50053b5 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -150,7 +150,7 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=Non from different sequences. Args: - tensor (Tensor): The input tensor to roll. + tensor (Tensor): The input tensor to roll. If None, returns (None, None). shifts (int): The shift of the tensor (typically -1 for MTP). dims (int): The dimension to roll (typically -1 for sequence dimension). cp_group (ProcessGroup): The context parallelism process group. If None or size=1, @@ -160,6 +160,9 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=Non Returns: tuple: (rolled_tensor, sum_of_rolled_tensor) """ + if tensor is None: + return None, None + # Handle packed sequences cases if packed_seq_params is not None: return _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group) @@ -631,6 +634,7 @@ def process_mtp_loss( cp_group: Optional[torch.distributed.ProcessGroup] = None, packed_seq_params: Optional[PackedSeqParams] = None, scale_logits_fn: Optional[Callable[[Tensor], Tensor]] = None, + input_ids: Optional[Tensor] = None, ) -> Tensor: """Process Multi-Token Prediction (MTP) loss computation. @@ -651,6 +655,10 @@ def process_mtp_loss( packed_seq_params (Optional[PackedSeqParams]): Packed sequence parameters. scale_logits_fn (Optional[Callable[[Tensor], Tensor]]): Optional function to scale logits before loss computation (e.g., MuP output scaling). + input_ids (Optional[Tensor]): Input token IDs. Used to derive labels when + ``labels`` is None (e.g. RL training), by rolling left to match the SFT + label convention (``label[i] = input_id[i + 1]``). Ignored when ``labels`` + is provided. Returns: Tensor: Updated hidden states after MTP loss processing (first chunk only). @@ -658,12 +666,27 @@ def process_mtp_loss( hidden_states_list = torch.chunk(hidden_states, 1 + config.mtp_num_layers, dim=0) hidden_states = hidden_states_list[0] + # When labels are not provided (e.g. RL training), derive them from input_ids by + # rolling left so that label[i] = input_id[i + 1], matching the SFT label format. + derived_labels_from_input_ids = False if labels is None: - return hidden_states + if input_ids is None: + return hidden_states + labels, _ = roll_tensor( + input_ids, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params + ) + derived_labels_from_input_ids = True mtp_labels = labels.clone() if loss_mask is None: loss_mask = torch.ones_like(mtp_labels) + if derived_labels_from_input_ids: + # input_ids has no real token beyond the sequence window, so the last rolled-in + # label is fabricated (zeroed). Roll loss_mask in lockstep with the + # input_ids -> labels shift so that boundary position is masked. + loss_mask, _ = roll_tensor( + loss_mask, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params + ) # Store the original number of tokens before rolling for proper normalization # when calculate_per_token_loss is enabled. This ensures MTP gradients are @@ -1095,6 +1118,7 @@ def _checkpointed_forward( hidden_states: Tensor, decoder_input: Tensor, attention_mask: Optional[Tensor] = None, + padding_mask: Optional[Tensor] = None, context: Optional[Tensor] = None, context_mask: Optional[Tensor] = None, rotary_pos_emb: Optional[Tensor] = None, @@ -1130,6 +1154,7 @@ def custom_forward( hidden_states, decoder_input, attention_mask, + padding_mask, context, context_mask, rotary_pos_emb, @@ -1141,6 +1166,7 @@ def custom_forward( hidden_states=hidden_states, decoder_input=decoder_input, attention_mask=attention_mask, + padding_mask=padding_mask, context=context, context_mask=context_mask, rotary_pos_emb=rotary_pos_emb, @@ -1187,6 +1213,7 @@ def checkpoint_handler(): hidden_states, decoder_input, attention_mask, + padding_mask, context, context_mask, rotary_pos_emb, @@ -1206,6 +1233,7 @@ def checkpoint_handler(): hidden_states, decoder_input, attention_mask, + padding_mask, context, context_mask, rotary_pos_emb, @@ -1232,6 +1260,7 @@ def checkpoint_handler(): hidden_states=hidden_states, decoder_input=decoder_input, attention_mask=attention_mask, + padding_mask=padding_mask, context=context, context_mask=context_mask, rotary_pos_emb=rotary_pos_emb, From dea749189255e96c1876214e19e7c510e522a3c7 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Mon, 6 Apr 2026 16:38:14 -0700 Subject: [PATCH 2/2] Add test for roll_tensor with None input Signed-off-by: Yi-Fu Wu Co-Authored-By: Claude Opus 4.6 (1M context) --- .../test_multi_token_prediction.py | 107 ++++++++++++++++++ 1 file changed, 107 insertions(+) diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index 583d4a16fa5..136c0b4d436 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -25,6 +25,7 @@ from megatron.core.transformer.multi_token_prediction import ( MTPLossLoggingHelper, MultiTokenPredictionBlock, + process_mtp_loss, roll_tensor, ) from megatron.core.transformer.transformer_config import TransformerConfig @@ -675,6 +676,112 @@ def test_packed_sequences_with_full_recompute(self): for name, param in gpt_model[0].named_parameters(): assert param.main_grad is not None, f"Gradient missing for {name}" + def test_roll_tensor_none_input(self): + """Test that roll_tensor returns (None, None) when given None input.""" + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=1) + result, sum_val = roll_tensor(None, shifts=-1, dims=-1) + assert result is None + assert sum_val is None + Utils.destroy_model_parallel() + + def test_roll_tensor_shifts_left_and_zeroes_last(self): + """Test that roll_tensor(-1) shifts left and zeroes the last position. + + This is the primitive used to derive MTP labels from input_ids when labels + are not provided (RL training): label[i] = input_id[i+1], last position zeroed. + The end-to-end derivation is covered by process_mtp_loss (see input_ids path). + """ + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=1) + # Simulate input_ids [batch=2, seq=5] + input_ids = torch.tensor( + [[10, 20, 30, 40, 50], [60, 70, 80, 90, 100]], dtype=torch.int64 + ).cuda() + rolled, _ = roll_tensor(input_ids, shifts=-1, dims=-1) + + # Expected: each row shifted left by 1, last element zeroed. + expected = torch.tensor( + [[20, 30, 40, 50, 0], [70, 80, 90, 100, 0]], dtype=torch.int64 + ).cuda() + assert torch.equal(rolled, expected) + Utils.destroy_model_parallel() + + def test_process_mtp_loss_skips_when_no_labels_and_no_input_ids(self): + """When labels and input_ids are both None, MTP loss is skipped (early return).""" + config = TransformerConfig( + hidden_size=8, num_layers=2, num_attention_heads=2, mtp_num_layers=1 + ) + hidden_states = torch.ones(2, 1, 4) + called = {'value': False} + + def output_layer(hidden, weight=None, runtime_gather_output=None): + return hidden.clone(), None + + def compute_language_model_loss(mtp_labels, mtp_logits): + called['value'] = True + return torch.ones_like(mtp_labels, dtype=mtp_logits.dtype) + + out = process_mtp_loss( + hidden_states=hidden_states, + labels=None, + loss_mask=None, + output_layer=output_layer, + output_weight=None, + runtime_gather_output=None, + is_training=False, + compute_language_model_loss=compute_language_model_loss, + config=config, + cp_group=None, + packed_seq_params=None, + input_ids=None, + ) + + # First chunk is returned unchanged and the loss is never computed. + assert not called['value'] + assert torch.equal(out, torch.chunk(hidden_states, 2, dim=0)[0]) + + def test_process_mtp_loss_derives_labels_from_input_ids(self): + """When labels is None (RL), labels are derived from input_ids by rolling left. + + process_mtp_loss rolls once to build the SFT-format labels (label[i] = + input_id[i+1]) and once more per MTP layer, so MTP head 0 targets input_id[i+2]. + The loss_mask is rolled in lockstep so the fabricated trailing label is masked. + """ + config = TransformerConfig( + hidden_size=8, num_layers=2, num_attention_heads=2, mtp_num_layers=1 + ) + # hidden_states is chunked into (1 + mtp_num_layers) along dim 0. + hidden_states = torch.ones(2, 1, 5) + input_ids = torch.tensor([[10, 20, 30, 40, 50]], dtype=torch.long) + seen = {'labels': None, 'masked_loss': None} + + def output_layer(hidden, weight=None, runtime_gather_output=None): + return hidden.clone(), None + + def compute_language_model_loss(mtp_labels, mtp_logits): + seen['labels'] = mtp_labels.clone() + # Per-position loss of 1.0 so loss_mask * loss exposes the active mask. + return torch.ones_like(mtp_labels, dtype=torch.float32) + + process_mtp_loss( + hidden_states=hidden_states, + labels=None, + loss_mask=None, + output_layer=output_layer, + output_weight=None, + runtime_gather_output=None, + is_training=False, + compute_language_model_loss=compute_language_model_loss, + config=config, + cp_group=None, + packed_seq_params=None, + input_ids=input_ids, + ) + + # input_ids rolled twice (once to SFT format, once in the MTP layer loop): + # [10,20,30,40,50] -> [20,30,40,50,0] -> [30,40,50,0,0]. + assert seen['labels'] is not None, "loss should be computed in RL mode" + assert torch.equal(seen['labels'], torch.tensor([[30, 40, 50, 0, 0]], dtype=torch.long)) + @pytest.mark.parametrize("cp", [1, 2]) def test_roll_tensor_with_packed_sequences(self, cp): """Test roll_tensor function with packed sequences, with and without CP.