diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index adccb559b27..4ceb7050f1d 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from __future__ import annotations import warnings @@ -640,7 +640,7 @@ def set_loss_scale(scale: torch.Tensor): def process_mtp_loss( hidden_states: Tensor, - labels: Tensor, + labels: Optional[Tensor], loss_mask: Optional[Tensor], output_layer: Callable, output_weight: Optional[Tensor], @@ -685,6 +685,23 @@ def process_mtp_loss( if loss_mask is None: loss_mask = torch.ones_like(mtp_labels) + output_weight_for_mtp = output_weight + output_layer_for_mtp = output_layer + if config.mtp_isolated_loss: + if output_weight_for_mtp is not None: + output_weight_for_mtp = output_weight_for_mtp.detach() + if isinstance(output_layer, torch.nn.Module): + output_layer_params = { + name: param.detach() for name, param in output_layer.named_parameters() + } + output_layer_buffers = dict(output_layer.named_buffers()) + output_layer_state = {**output_layer_params, **output_layer_buffers} + + def output_layer_for_mtp(input_: Tensor, **kwargs): + return torch.func.functional_call( + output_layer, output_layer_state, args=(input_,), kwargs=kwargs + ) + # Store the original number of tokens before rolling for proper normalization # when calculate_per_token_loss is enabled. This ensures MTP gradients are # correctly scaled relative to the main loss gradients in finalize_model_grads. @@ -701,17 +718,17 @@ def process_mtp_loss( loss_mask, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params ) if fuse_linear_cross_entropy: - mtp_loss = output_layer( + mtp_loss = output_layer_for_mtp( hidden_states_list[mtp_layer_number + 1], - weight=output_weight, + weight=output_weight_for_mtp, runtime_gather_output=runtime_gather_output, output_cross_entropy_loss=True, labels=mtp_labels, ) else: - mtp_logits, _ = output_layer( + mtp_logits, _ = output_layer_for_mtp( hidden_states_list[mtp_layer_number + 1], - weight=output_weight, + weight=output_weight_for_mtp, runtime_gather_output=runtime_gather_output, ) if scale_logits_fn is not None: @@ -991,6 +1008,8 @@ def _get_embeddings( ) # embedding decoder_input = embedding(input_ids=input_ids, position_ids=position_ids) + if self.config.mtp_isolated_loss: + decoder_input = decoder_input.detach() hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) @@ -1724,6 +1743,11 @@ def forward( hidden_states = mhc_chunks[offset] else: hidden_states = hidden_states_list[offset] + if self.config.mtp_isolated_loss: + hidden_states = hidden_states.detach().requires_grad_(True) + hidden_states = make_viewless_tensor( + inp=hidden_states, requires_grad=True, keep_graph=False + ) for iteration in range(self.config.mtp_num_layers): layer_idx = 0 if self.mtp_use_repeated_layer else iteration (hidden_states, input_ids, position_ids) = self.layers[layer_idx]( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 775f19e7aeb..bd9777087bb 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -81,6 +81,15 @@ class TransformerConfig(ModelParallelConfig): which serves as an additional training objective. """ + mtp_isolated_loss: bool = False + """If True, MTP loss only updates MTP module parameters. The MTP loss graph is + detached from the main decoder, shared embeddings, and output layer weights. + + For online RL, keep ``labels=None`` so the main LM head returns logits for the + external RL loss. MTP auxiliary loss can still be trained by deriving its labels + from ``input_ids`` in the MTP loss path; this option isolates that auxiliary loss + from the main model parameters.""" + mtp_use_repeated_layer: bool = False """Use a single MTP layer repeatedly instead of multiple separate layers.""" diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index a84389a8057..72e197b5f1a 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -227,6 +227,7 @@ "mup_output_mult": 1.0, "mup_width_mult": 1.0, "mtp_hybrid_override_pattern": None, + "mtp_isolated_loss": False, "mtp_loss_scaling_factor": 0.1, "mtp_num_layers": None, "mtp_standalone": False, diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index 580fd23d783..a225cd376a6 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -1,10 +1,12 @@ -# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import os import sys +from types import SimpleNamespace import pytest import torch +import torch.nn.functional as F from torch import Tensor from megatron.core.enums import ModelType @@ -26,6 +28,7 @@ from megatron.core.transformer.multi_token_prediction import ( MTPLossLoggingHelper, MultiTokenPredictionBlock, + process_mtp_loss, roll_tensor, ) from megatron.core.transformer.transformer_block import TransformerBlock @@ -52,6 +55,40 @@ _SEED = 42 +class _TestOutputLayer(torch.nn.Module): + def __init__(self, hidden_size, vocab_size): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(vocab_size, hidden_size)) + + def forward( + self, + input_, + weight=None, + runtime_gather_output=None, + output_cross_entropy_loss=False, + labels=None, + ): + del runtime_gather_output + weight = self.weight if weight is None else weight + logits = torch.matmul(input_, weight.t()) + if output_cross_entropy_loss: + logits = logits.transpose(0, 1).contiguous() + loss = F.cross_entropy( + logits.view(-1, logits.size(-1)), labels.reshape(-1), reduction='none' + ) + return loss.view_as(labels) + return logits, None + + +class _ScaleMTPLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.scale = torch.nn.Parameter(torch.tensor(2.0)) + + def forward(self, input_ids, position_ids, hidden_states, **_kwargs): + return hidden_states * self.scale, input_ids, position_ids + + class TestMultiTokenPredictionLayer: def setup_method(self, method): os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' @@ -134,6 +171,105 @@ def test_constructor_ues_te(self, tp, cp): assert num_weights == 15216 * config.mtp_num_layers +class TestProcessMTPLoss: + def test_isolated_loss_detaches_encoder_hidden_states(self): + """MTP isolated loss should not update the main decoder hidden states.""" + + torch.manual_seed(_SEED) + seq_length = 4 + micro_batch_size = 2 + hidden_size = 8 + input_ids = torch.arange(seq_length).repeat(micro_batch_size, 1) + position_ids = torch.arange(seq_length).repeat(micro_batch_size, 1) + + for isolated_loss in (False, True): + config = TransformerConfig( + num_layers=1, + hidden_size=hidden_size, + num_attention_heads=1, + mtp_num_layers=1, + mtp_loss_scaling_factor=1.0, + mtp_isolated_loss=isolated_loss, + ) + mtp_layer = _ScaleMTPLayer() + mtp_block = SimpleNamespace( + config=config, vp_stage=None, mtp_use_repeated_layer=False, layers=[mtp_layer] + ) + hidden_states = torch.randn( + seq_length, micro_batch_size, hidden_size, requires_grad=True + ) + + output = MultiTokenPredictionBlock.forward( + mtp_block, + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + attention_mask=None, + ) + mtp_output = output[seq_length:] + mtp_output.sum().backward() + + assert mtp_layer.scale.grad is not None + if isolated_loss: + assert hidden_states.grad is None or torch.count_nonzero(hidden_states.grad) == 0 + else: + assert hidden_states.grad is not None + assert torch.count_nonzero(hidden_states.grad) > 0 + + def test_isolated_loss_detaches_output_layer(self): + """MTP isolated loss should not update output layer weights.""" + + def compute_language_model_loss(labels, logits): + logits = logits.transpose(0, 1).contiguous() + loss = F.cross_entropy( + logits.view(-1, logits.size(-1)), labels.reshape(-1), reduction='none' + ) + return loss.view_as(labels) + + torch.manual_seed(_SEED) + seq_length = 4 + micro_batch_size = 2 + hidden_size = 8 + vocab_size = 16 + labels = torch.randint(vocab_size, (micro_batch_size, seq_length)) + loss_mask = torch.ones_like(labels, dtype=torch.float32) + + for isolated_loss in (False, True): + config = TransformerConfig( + num_layers=1, + hidden_size=hidden_size, + num_attention_heads=1, + mtp_num_layers=1, + mtp_loss_scaling_factor=1.0, + mtp_isolated_loss=isolated_loss, + ) + output_layer = _TestOutputLayer(hidden_size, vocab_size) + hidden_states = torch.randn( + seq_length * (1 + config.mtp_num_layers), + micro_batch_size, + hidden_size, + requires_grad=True, + ) + + output = process_mtp_loss( + hidden_states=hidden_states, + labels=labels, + loss_mask=loss_mask, + output_layer=output_layer, + output_weight=None, + runtime_gather_output=False, + is_training=False, + compute_language_model_loss=compute_language_model_loss, + config=config, + ) + output.sum().backward() + + if isolated_loss: + assert output_layer.weight.grad is None + else: + assert output_layer.weight.grad is not None + + class TestMultiTokenPrediction: def setup_method(self, method): self.seq_length = 32