diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index d9d22587a07..ee23af3496a 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -620,7 +620,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], @@ -665,15 +665,32 @@ 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. original_num_tokens = loss_mask.sum() for mtp_layer_number in range(config.mtp_num_layers): - 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: @@ -927,6 +944,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) @@ -1614,6 +1633,11 @@ def forward( offset = get_mtp_layer_offset(self.config, self.vp_stage) hidden_states_list = list(torch.chunk(hidden_states, 1 + offset, dim=0)) 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, padding_mask) = self.layers[layer_idx]( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 95150103f27..d4a91a2a46d 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -78,6 +78,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 5dce465cfe6..b0d821d898a 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -217,6 +217,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 583d4a16fa5..73ff84e97a9 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -3,9 +3,11 @@ import os import sys import types +from types import SimpleNamespace import pytest import torch +import torch.nn.functional as F from megatron.core.enums import ModelType from megatron.core.extensions.transformer_engine import HAVE_TE @@ -25,6 +27,7 @@ from megatron.core.transformer.multi_token_prediction import ( MTPLossLoggingHelper, MultiTokenPredictionBlock, + process_mtp_loss, roll_tensor, ) from megatron.core.transformer.transformer_config import TransformerConfig @@ -49,6 +52,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, padding_mask=None, **_kwargs): + return hidden_states * self.scale, input_ids, position_ids, padding_mask + + class TestMultiTokenPredictionLayer: def setup_method(self, method): os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' @@ -226,6 +263,105 @@ def fake_proj_and_transformer_layer( assert torch.equal(returned_padding_mask, expected_padding_mask) +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