Skip to content

Support isolated MTP loss #5098

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Victarry wants to merge 1 commit into from
Closed

Support isolated MTP loss #5098

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 27 additions & 3 deletions megatron/core/transformer/multi_token_prediction.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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](
Expand Down
9 changes: 9 additions & 0 deletions megatron/core/transformer/transformer_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand Down
1 change: 1 addition & 0 deletions tests/unit_tests/models/test_hybrid_moe_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
136 changes: 136 additions & 0 deletions tests/unit_tests/transformer/test_multi_token_prediction.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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'
Expand Down Expand Up @@ -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
Expand Down