From e8505915dbf7e808cc1f425537e23e41407f84d4 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Thu, 15 Oct 2026 03:44:51 +0000 Subject: [PATCH] [Fix] Isolate generic MTP training and future cache fixes Co-authored-by: Kuang <128125671+Knlife@users.noreply.github.com> --- tests/loss/test_mtp_e2e_tv_loss.py | 23 ++ tests/loss/test_mtp_empty_shards.py | 98 ++++++ .../model/test_mtp_local_future_embeddings.py | 321 ++++++++++++++++++ xtuner/v1/loss/mtp_loss.py | 6 +- xtuner/v1/module/mtp/config.py | 6 + xtuner/v1/module/mtp/mtp_block.py | 52 ++- xtuner/v1/module/mtp/utils.py | 37 ++ 7 files changed, 535 insertions(+), 8 deletions(-) create mode 100644 tests/loss/test_mtp_empty_shards.py create mode 100644 tests/model/test_mtp_local_future_embeddings.py diff --git a/tests/loss/test_mtp_e2e_tv_loss.py b/tests/loss/test_mtp_e2e_tv_loss.py index f0f7c043f6..4861ed8b41 100644 --- a/tests/loss/test_mtp_e2e_tv_loss.py +++ b/tests/loss/test_mtp_e2e_tv_loss.py @@ -1,5 +1,6 @@ from types import SimpleNamespace +import pytest import torch import torch.nn.functional as F @@ -72,3 +73,25 @@ def test_e2e_tv_can_detach_shared_lm_head(): assert target.grad is None assert head_weight.grad is None assert all(draft.grad is not None and torch.count_nonzero(draft.grad) > 0 for draft in drafts) + + +@pytest.mark.parametrize("detach_head", [False, True]) +def test_e2e_tv_empty_bf16_shard_matches_nonempty_loss_dtype(detach_head): + torch.manual_seed(2) + target = torch.randn(1, 8, 4, device=DEVICE, dtype=torch.bfloat16, requires_grad=True) + drafts = [torch.randn_like(target, requires_grad=True) for _ in range(2)] + head_weight = torch.randn(7, 4, device=DEVICE, dtype=torch.bfloat16, requires_grad=True) + ctx = _build_context(detach_head=detach_head) + nonempty_loss, _ = ctx.forward((target, drafts), head_weight) + ctx.loss_kwargs.loss_weight.zero_() + empty_loss, _ = ctx.forward((target, drafts), head_weight) + + assert empty_loss.dtype == nonempty_loss.dtype == torch.float32 + assert empty_loss.item() == 0 + empty_loss.backward() + assert target.grad is None + if detach_head: + assert head_weight.grad is None + else: + assert head_weight.grad is not None and torch.count_nonzero(head_weight.grad) == 0 + assert all(draft.grad is not None and torch.count_nonzero(draft.grad) == 0 for draft in drafts) diff --git a/tests/loss/test_mtp_empty_shards.py b/tests/loss/test_mtp_empty_shards.py new file mode 100644 index 0000000000..10a9d1d0f5 --- /dev/null +++ b/tests/loss/test_mtp_empty_shards.py @@ -0,0 +1,98 @@ +"""BF16 MTP CE/TV losses with a completely unsupervised SP shard.""" + +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.distributed.device_mesh import init_device_mesh + +from xtuner.v1.loss import mtp_loss +from xtuner.v1.loss.mtp_loss import MTPE2ETVLossConfig, MTPE2ETVLossContext, MTPLossConfig, MTPLossContext + + +@pytest.fixture(autouse=True) +def _cpu_loss_device(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(mtp_loss, "DEVICE", torch.device("cpu")) + + +@pytest.mark.parametrize("mode", ["eager", "chunk"]) +@pytest.mark.parametrize("detach_head", [False, True]) +def test_mtp_ce_empty_bf16_shard_has_fp32_loss_and_zero_gradients(mode: str, detach_head: bool): + torch.manual_seed(0) + hidden = torch.randn(1, 18, 4, dtype=torch.bfloat16, requires_grad=True) + head = torch.randn(7, 4, dtype=torch.bfloat16, requires_grad=True) + seq_ctx = SimpleNamespace(cu_seq_lens_k=torch.tensor([0, 10, 18], dtype=torch.int32)) + cfg = MTPLossConfig(mode=mode, chunk_size=2, mtp_depth=7, detach_mtp_lm_head_weight=detach_head) + ctx = cfg.build({"shifted_labels": torch.arange(18).remainder(7).unsqueeze(0), "seq_ctx": seq_ctx}) + assert ctx is not None + ctx = MTPLossContext.build_batches([ctx], cu_seq_lens_list=[seq_ctx.cu_seq_lens_k])[0] + nonempty_loss, _ = ctx.forward(hidden, head) + ctx.loss_kwargs.shifted_labels.fill_(-100) + ctx.loss_kwargs.loss_weight.zero_() + empty_loss, _ = ctx.forward(hidden, head) + + assert empty_loss.dtype == nonempty_loss.dtype == torch.float32 + assert empty_loss.item() == 0 + empty_loss.backward() + assert hidden.grad is not None and torch.count_nonzero(hidden.grad) == 0 + if detach_head: + assert head.grad is None + else: + assert head.grad is not None and torch.count_nonzero(head.grad) == 0 + + +def _mixed_sp_worker(rank: int, rendezvous: str) -> None: + mtp_loss.DEVICE = torch.device("cpu") + torch.set_num_threads(1) + dist.init_process_group( + "gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=2, timeout=timedelta(seconds=60) + ) + try: + mesh = init_device_mesh("cpu", (2,), mesh_dim_names=("sp",)) + labels = torch.cat([torch.arange(10).remainder(7), torch.full((10,), -100)]).unsqueeze(0) + seq_ctx = SimpleNamespace(cu_seq_lens_k=torch.tensor([0, 20], dtype=torch.int32)) + data = {"shifted_labels": labels, "seq_ctx": seq_ctx} + for objective in ("ce", "e2e_tv"): + torch.manual_seed(1) + target = torch.randn(1, 10, 4, dtype=torch.bfloat16, requires_grad=True) + drafts = [torch.randn_like(target, requires_grad=True) for _ in range(7)] + head = torch.randn(7, 4, dtype=torch.bfloat16, requires_grad=True) + if objective == "ce": + losses = [] + for depth, draft in enumerate(drafts, start=1): + cfg = MTPLossConfig(mode="chunk", chunk_size=2, mtp_depth=depth) + ctx = cfg.build(data, sp_mesh=mesh) + assert ctx is not None + ctx = MTPLossContext.build_batches([ctx], [seq_ctx.cu_seq_lens_k], sp_mesh=mesh)[0] + if rank == 1: + assert torch.count_nonzero(ctx.loss_kwargs.loss_weight) == 0 + losses.append(ctx.forward(draft, head)[0]) + loss = torch.stack(losses).mean() * 0.1 + else: + cfg = MTPE2ETVLossConfig(mode="chunk", chunk_size=2, num_steps=7) + ctx = cfg.build(data, sp_mesh=mesh) + assert ctx is not None + ctx = MTPE2ETVLossContext.build_batches([ctx], [seq_ctx.cu_seq_lens_k], sp_mesh=mesh)[0] + if rank == 1: + assert torch.count_nonzero(ctx.loss_kwargs.loss_weight) == 0 + loss = ctx.forward((target, drafts), head)[0] * 0.1 + + assert loss.dtype == torch.float32 and torch.isfinite(loss) and loss.item() > 0 + gathered = [torch.empty_like(loss) for _ in range(2)] + dist.all_gather(gathered, loss.detach()) + torch.testing.assert_close(gathered[0], gathered[1], rtol=0, atol=0) + loss.backward() + assert target.grad is None + for tensor in [head, *drafts]: + assert tensor.grad is not None and torch.isfinite(tensor.grad).all() + assert (torch.count_nonzero(tensor.grad) == 0).item() == (rank == 1) + finally: + dist.destroy_process_group() + + +def test_mtp_seven_depth_ce_and_tv_agree_across_mixed_empty_sp_ranks(tmp_path: Path): + mp.spawn(_mixed_sp_worker, args=(str(tmp_path / "mixed-mtp-shards"),), nprocs=2, join=True) diff --git a/tests/model/test_mtp_local_future_embeddings.py b/tests/model/test_mtp_local_future_embeddings.py new file mode 100644 index 0000000000..a158e26395 --- /dev/null +++ b/tests/model/test_mtp_local_future_embeddings.py @@ -0,0 +1,321 @@ +"""Local MTP buffers must preserve packed rolling values, gradients and lifetimes.""" + +from datetime import timedelta +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch import nn +from torch.distributed.device_mesh import init_device_mesh +from torch.utils.checkpoint import checkpoint + +from xtuner.v1.data_proto import SequenceContext +from xtuner.v1.loss.mtp_loss import MTPE2ETVLossConfig, MTPE2ETVLossContext, MTPE2ETVLossKwargs +from xtuner.v1.module.mtp import MTPBlock, MTPConfig +from xtuner.v1.module.mtp.utils import local_packed_future_embeddings, roll_packed_tensor + + +BOUNDARIES = torch.tensor([0, 2, 2, 9, 10, 17, 20], dtype=torch.int32) + + +@pytest.mark.parametrize("depth", range(1, 8)) +@pytest.mark.parametrize("shard_start", [0, 8, 16]) +@pytest.mark.parametrize("dtype", [torch.float64, torch.bfloat16]) +@pytest.mark.parametrize( + "device", + ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA"))], +) +def test_local_future_matches_full_roll_at_pack_and_sp_boundaries(depth, shard_start, dtype, device): + source = torch.arange(60, dtype=dtype, device=device).reshape(1, 20, 3).requires_grad_() + reference_source = source.detach().clone().requires_grad_() + boundaries = BOUNDARIES.to(device) + local = local_packed_future_embeddings(source, boundaries, depth=depth, shard_start=shard_start, shard_size=4) + reference = roll_packed_tensor(reference_source, boundaries, shifts=-depth, dim=1) + torch.testing.assert_close(local, reference[:, shard_start : shard_start + 4], rtol=0, atol=0) + assert local.is_contiguous() and local.storage_offset() == 0 + assert local.untyped_storage().nbytes() == local.numel() * local.element_size() + assert local.untyped_storage().data_ptr() != source.untyped_storage().data_ptr() + local.float().square().sum().backward() + reference[:, shard_start : shard_start + 4].float().square().sum().backward() + torch.testing.assert_close(source.grad, reference_source.grad, rtol=0, atol=0) + + +class _Mesh: + def __init__(self, rank): + self.rank = rank + + def size(self): + return 5 + + def get_local_rank(self): + return self.rank + + def get_group(self): + return None + + +def _context(embeddings, rank=None): + start = 0 if rank is None else rank * 4 + size = 20 if rank is None else 4 + ids = torch.arange(20).unsqueeze(0) + return SequenceContext( + input_ids=ids[:, start : start + size], + inputs_embeds=embeddings, + raw_input_ids=ids, + cu_seq_lens_q=BOUNDARIES, + cu_seq_lens_k=BOUNDARIES, + max_length_q=7, + max_length_k=7, + num_padding=3 if rank in (None, 4) else 0, + position_ids=ids[:, start : start + size], + sequence_parallel_mesh=None if rank is None else _Mesh(rank), + shard_start=start, + shard_size=size, + ) + + +class _CheckpointLayer(nn.Module): + def __init__(self, recompute, reentrant=False): + super().__init__() + self.weight = nn.Parameter(torch.tensor(0.5, dtype=torch.float64)) + self.recompute = recompute + self.reentrant = reentrant + self.futures = [] + self.contexts = [] + + def _compute(self, hidden, future, ctx): + return (self.weight * hidden + future).tanh() * ctx.mask.unsqueeze(-1) + + def forward(self, hidden_states, *, future_embeddings, position_embeddings, seq_ctx): + multi = isinstance(hidden_states, list) + states = hidden_states if multi else [hidden_states] + futures = future_embeddings if multi else [future_embeddings] + contexts = seq_ctx if multi else [seq_ctx] + self.futures.append(futures) + self.contexts.append(contexts) + outputs = [ + checkpoint(self._compute, h, e, ctx, use_reentrant=self.reentrant) + if self.recompute + else self._compute(h, e, ctx) + for h, e, ctx in zip(states, futures, contexts) + ] + result = dict(hidden_states=outputs if multi else outputs[0]) + result.update( + { + key: [None] * len(states) if multi else None + for key in ("router_logits", "router_weights", "router_topk_ids") + } + ) + return result + + +def _run(optimized, recompute, micro_batches, rank=None, detach=False): + layer = _CheckpointLayer(recompute) + cache_config = {} if optimized is None else {"use_local_future_embeddings": optimized} + block = MTPBlock( + mtp_config=MTPConfig(num_layers=7, share_weights=True, detach_mtp_inputs=detach, **cache_config), + mtp_layers=[layer], + ) + # Values stay small so every prediction depth contributes measurable gradients. + full = torch.arange(60, dtype=torch.float64).reshape(1, 20, 3) / 100 + source = ( + full.clone().requires_grad_() if rank is None else full[:, rank * 4 : rank * 4 + 4].clone().requires_grad_() + ) + contexts = [_context(source, rank) for _ in range(micro_batches)] + states = [torch.full_like(source, 0.2 + mb / 10, requires_grad=True) for mb in range(micro_batches)] + positions = [(torch.empty(0), torch.empty(0)) for _ in contexts] + with patch( + "xtuner.v1.data_proto.sequence_context.all_gather_tensor_autograd", + return_value=full.detach().clone(), + ) as gather: + outputs = block( + states[0] if micro_batches == 1 else states, + embed_tokens_fn=lambda _: pytest.fail("Supplied embeddings must not be re-embedded"), + position_embeddings=positions[0] if micro_batches == 1 else positions, + seq_ctx=contexts[0] if micro_batches == 1 else contexts, + ) + groups = [outputs] if micro_batches == 1 else outputs + sum(out["hidden_states"].square().mean() for group in groups for out in group).backward() + return layer, source, states, groups, gather.call_count, contexts + + +@pytest.mark.parametrize("recompute", [False, True]) +@pytest.mark.parametrize("micro_batches", [1, 2]) +@pytest.mark.parametrize("rank", [None, 0, 4]) +def test_seven_depths_preserve_outputs_and_backward_with_checkpoint(recompute, micro_batches, rank): + old_layer, old_source, old_states, old_outputs, old_gathers, _ = _run(False, recompute, micro_batches, rank) + layer, source, states, outputs, gathers, contexts = _run(True, recompute, micro_batches, rank) + assert gathers == old_gathers == (0 if rank is None else micro_batches) + torch.testing.assert_close(layer.weight.grad, old_layer.weight.grad) + for state, old_state, group, old_group in zip(states, old_states, outputs, old_outputs): + torch.testing.assert_close(state.grad, old_state.grad) + for out, old_out in zip(group, old_group): + torch.testing.assert_close(out["hidden_states"], old_out["hidden_states"]) + if rank is None: + torch.testing.assert_close(source.grad, old_source.grad) + else: + # This fake collective returns a detached tensor; the distributed cases below + # check gradients through the real autograd-aware collective. + assert source.grad is old_source.grad is None + for depth in range(7): + for mb in range(micro_batches): + assert layer.contexts[depth][mb] is contexts[mb] + future = layer.futures[depth][mb] + torch.testing.assert_close(future, old_layer.futures[depth][mb], rtol=0, atol=0) + assert future.is_contiguous() and future.storage_offset() == 0 + assert future.untyped_storage().nbytes() == future.numel() * future.element_size() + # Hold all seven inputs across backward; each is an independent local allocation. + for mb in range(micro_batches): + pointers = {futures[mb].untyped_storage().data_ptr() for futures in layer.futures} + assert len(pointers) == 7 + + +def test_local_cache_preserves_single_microbatch_input_detach(): + old_layer, old_source, old_states, _, _, _ = _run(False, True, 1, detach=True) + layer, source, states, _, _, _ = _run(True, True, 1, detach=True) + assert source.grad is old_source.grad is states[0].grad is old_states[0].grad is None + torch.testing.assert_close(layer.weight.grad, old_layer.weight.grad) + + +def test_default_uses_local_future_buffers_and_retains_explicit_legacy_opt_out(): + assert MTPConfig(num_layers=7).use_local_future_embeddings is True + legacy = MTPConfig(num_layers=7, use_local_future_embeddings=False) + assert MTPConfig.model_validate_json(legacy.model_dump_json()).use_local_future_embeddings is False + layer, source, _, _, _, contexts = _run(None, True, 1, rank=4) + assert all(depth_contexts[0] is contexts[0] for depth_contexts in layer.contexts) + assert all( + futures[0].untyped_storage().nbytes() == source.numel() * source.element_size() for futures in layer.futures + ) + + +def test_token_id_only_inputs_keep_legacy_rolling_when_optimization_is_enabled(): + ctx = SequenceContext.from_input_ids((torch.arange(20).unsqueeze(0),), device="cpu") + layer = _CheckpointLayer(False) + block = MTPBlock( + mtp_config=MTPConfig(num_layers=7, share_weights=True), + mtp_layers=[layer], + ) + block( + torch.zeros(1, 20, 3), + embed_tokens_fn=lambda ids: ids.double().unsqueeze(-1).expand(-1, -1, 3), + position_embeddings=(torch.empty(0), torch.empty(0)), + seq_ctx=ctx, + ) + for depth, contexts in enumerate(layer.contexts, start=1): + expected = torch.arange(20).roll(-depth) + expected[-depth:] = 0 + torch.testing.assert_close(contexts[0].input_ids, expected.unsqueeze(0)) + + +def _distributed_cache_worker(rank, rendezvous, backend="gloo", reentrant=False): + torch.set_num_threads(1) + cuda = backend == "nccl" + device = torch.device(f"cuda:{rank}" if cuda else "cpu") + dtype = torch.bfloat16 if cuda else torch.float64 + if cuda: + torch.cuda.set_device(device) + dist.init_process_group( + backend, init_method=f"file://{rendezvous}", rank=rank, world_size=2, timeout=timedelta(seconds=120) + ) + try: + mesh = init_device_mesh(device.type, (2,), mesh_dim_names=("sp",)) + # NCCL includes a long enough pack for seven-depth TV supervision on both ranks. + pack_ends = [0, 2, 2, 7, 20] if cuda else [0, 2, 2, 7, 13, 17, 20] + boundaries = torch.tensor(pack_ends, dtype=torch.int32, device=device) + full = torch.arange(60, dtype=dtype, device=device).reshape(1, 20, 3) / 100 + reference = None + for optimized in (False, True): + layer = _CheckpointLayer(True, reentrant=reentrant).to(device=device, dtype=dtype) + block = MTPBlock( + mtp_config=MTPConfig(num_layers=7, share_weights=True, use_local_future_embeddings=optimized), + mtp_layers=[layer], + ) + source = full[:, rank * 10 : rank * 10 + 10].clone().requires_grad_() + state = torch.full_like(source, 0.2, requires_grad=True) + head = (torch.arange(39, device=device, dtype=dtype).reshape(13, 3) / 100).requires_grad_() + ctx = SequenceContext( + input_ids=None, + inputs_embeds=source, + raw_input_ids=torch.arange(20, device=device).unsqueeze(0), + cu_seq_lens_q=boundaries, + cu_seq_lens_k=boundaries, + max_length_q=13 if cuda else 6, + max_length_k=13 if cuda else 6, + sequence_parallel_mesh=mesh, + position_ids=torch.arange(rank * 10, rank * 10 + 10, device=device).unsqueeze(0), + shard_start=rank * 10, + shard_size=10, + device=device, + ) + + outputs = block( + state, + embed_tokens_fn=lambda _: None, + position_embeddings=(torch.empty(0, device=device), torch.empty(0, device=device)), + seq_ctx=ctx, + ) + if cuda: + positions = torch.arange(rank * 10, rank * 10 + 10, device=device) + ends = boundaries[1:][torch.searchsorted(boundaries[1:], positions, right=True)] + weights = ((positions + 7) < ends).float().unsqueeze(0) / 6 + tv_ctx = MTPE2ETVLossContext( + MTPE2ETVLossConfig(num_steps=7, mode="chunk", chunk_size=4), + MTPE2ETVLossKwargs( + shifted_labels=torch.zeros(1, 10, dtype=torch.long, device=device), + loss_weight=weights, + sp_mesh=mesh, + ), + ) + loss, _ = tv_ctx.forward((state, [out["hidden_states"] for out in outputs]), head) + assert torch.isfinite(loss) and loss.item() > 0 + else: + loss = sum(out["hidden_states"].square().mean() for out in outputs) + loss.backward() + assert source.grad is not None and source.grad.isfinite().all() and state.grad is not None + for depth, futures in enumerate(layer.futures, start=1): + expected = roll_packed_tensor(full, boundaries, shifts=-depth, dim=1) + torch.testing.assert_close(futures[0], expected[:, rank * 10 : rank * 10 + 10], rtol=0, atol=0) + if reference is None: + reference = ( + layer.weight.grad.clone(), + state.grad.clone(), + [out["hidden_states"] for out in outputs], + loss.detach().clone(), + head.grad.clone() if cuda else None, + source.grad.clone(), + ) + else: + torch.testing.assert_close(loss.detach(), reference[3], rtol=0, atol=0) + if cuda: + assert head.grad is not None and head.grad.isfinite().all() and head.grad.abs().sum() > 0 + torch.testing.assert_close(head.grad, reference[4], rtol=0, atol=0) + torch.testing.assert_close(layer.weight.grad, reference[0], rtol=0, atol=0) + torch.testing.assert_close(state.grad, reference[1], rtol=0, atol=0) + torch.testing.assert_close(source.grad, reference[5], rtol=0, atol=0) + for out, old_out in zip(outputs, reference[2]): + torch.testing.assert_close(out["hidden_states"], old_out, rtol=0, atol=0) + assert all(contexts[0] is ctx for contexts in layer.contexts) + assert all( + futures[0].untyped_storage().nbytes() == source.numel() * source.element_size() + for futures in layer.futures + ) + finally: + dist.destroy_process_group() + + +def test_distributed_local_future_cache_matches_full_roll_and_checkpoint_backward(tmp_path): + mp.spawn(_distributed_cache_worker, args=(str(tmp_path / "local-future-cache"),), nprocs=2, join=True) + + +@pytest.mark.gpu +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices for NCCL SP") +@pytest.mark.parametrize("reentrant", [False, True]) +def test_nccl_local_future_cache_matches_full_roll_and_checkpoint_backward(tmp_path, reentrant): + mp.spawn( + _distributed_cache_worker, + args=(str(tmp_path / "nccl-local-future-cache"), "nccl", reentrant), + nprocs=2, + join=True, + ) diff --git a/xtuner/v1/loss/mtp_loss.py b/xtuner/v1/loss/mtp_loss.py index ad1aff5f18..dc374bcfe2 100644 --- a/xtuner/v1/loss/mtp_loss.py +++ b/xtuner/v1/loss/mtp_loss.py @@ -398,7 +398,11 @@ def forward( # type: ignore[override] valid_mask = loss_weight != 0 if not valid_mask.any(): - loss = head_weight.sum() * 0.0 + sum(draft.sum() * 0.0 for draft in draft_hidden_states) + # Empty SP shards must match the FP32 loss on non-empty ranks + # before the collective, while retaining zero-gradient edges. + loss = head_weight.sum(dtype=torch.float32) * 0.0 + sum( + draft.sum(dtype=torch.float32) * 0.0 for draft in draft_hidden_states + ) else: selected_weight = loss_weight[valid_mask] alpha_per_step: list[torch.Tensor] = [] diff --git a/xtuner/v1/module/mtp/config.py b/xtuner/v1/module/mtp/config.py index 84b390d6b9..df1b3f83ea 100644 --- a/xtuner/v1/module/mtp/config.py +++ b/xtuner/v1/module/mtp/config.py @@ -36,6 +36,11 @@ class MTPConfig(BaseModel): expected multi-step rejection-sampling acceptance length. Default: ``"ce"``. tv_loss_chunk_size (int): Number of token positions processed together when computing the exact full-vocabulary TV overlap. Default: 128. + use_local_future_embeddings (bool): Gather supplied embeddings once, then + materialize only the local SP slice for each prediction depth. All depths + reuse the unshifted sequence metadata; future tokens are passed explicitly + via ``future_embeddings``. Token-ID-only inputs keep the rolling path. + Default: True. Set False to restore the full-sequence rolling path. Example: >>> # In model config @@ -58,3 +63,4 @@ class MTPConfig(BaseModel): loss_scaling_factor: Annotated[float, Parameter(group="model")] = 0.1 loss_type: Annotated[Literal["ce", "e2e_tv"], Parameter(group="model")] = "ce" tv_loss_chunk_size: Annotated[int, Parameter(group="model")] = 128 + use_local_future_embeddings: Annotated[bool, Parameter(group="model")] = True diff --git a/xtuner/v1/module/mtp/mtp_block.py b/xtuner/v1/module/mtp/mtp_block.py index f13f1c2235..4957c30fb9 100644 --- a/xtuner/v1/module/mtp/mtp_block.py +++ b/xtuner/v1/module/mtp/mtp_block.py @@ -13,7 +13,7 @@ from .config import MTPConfig from .mtp_layer import MTPLayer -from .utils import roll_sequence_context +from .utils import local_packed_future_embeddings, roll_sequence_context MTPDepthOutput = MoEDecoderLayerOutput @@ -177,15 +177,15 @@ def _forward( mtp_outputs: list[MTPDepthOutput] = [] current_hidden_states = hidden_states.detach() if self.mtp_config.detach_mtp_inputs else hidden_states current_seq_ctx = seq_ctx + raw_inputs_embeds = self._gather_future_embeddings(seq_ctx) previous_layer_results: MTPInternalOutput | None = None num_steps = self.mtp_config.num_layers for step in range(num_steps): layer = cast(MTPLayer, self.layers[0] if self.mtp_config.share_weights else self.layers[step]) - # Roll each packed sequence independently so we get the (i+k)-th token while - # respecting per-sequence boundaries inside the packed batch. - current_seq_ctx = roll_sequence_context(current_seq_ctx, shifts=-1) - future_embeddings = self._embed_future(current_seq_ctx, embed_tokens_fn) + current_seq_ctx, future_embeddings = self._next_future_inputs( + current_seq_ctx, embed_tokens_fn, depth=step + 1, raw_inputs_embeds=raw_inputs_embeds + ) if self.mtp_config.detach_mtp_inputs: future_embeddings = future_embeddings.detach() @@ -250,14 +250,19 @@ def _micro_batch_forward( outputs_per_mb: list[list[MTPDepthOutput]] = [[] for _ in range(n)] current_hidden_states_list = list(hidden_states_list) current_seq_ctx_list = list(seq_ctx_list) + raw_inputs_embeds_list = [self._gather_future_embeddings(ctx) for ctx in seq_ctx_list] previous_layer_results: MTPInternalOutput | None = None num_steps = self.mtp_config.num_layers for step in range(num_steps): layer = cast(MTPLayer, self.layers[0] if self.mtp_config.share_weights else self.layers[step]) - current_seq_ctx_list = [roll_sequence_context(ctx, shifts=-1) for ctx in current_seq_ctx_list] - future_embeddings_list = [self._embed_future(ctx, embed_tokens_fn) for ctx in current_seq_ctx_list] + future_inputs = [ + self._next_future_inputs(ctx, embed_tokens_fn, depth=step + 1, raw_inputs_embeds=raw) + for ctx, raw in zip(current_seq_ctx_list, raw_inputs_embeds_list) + ] + current_seq_ctx_list = [ctx for ctx, _ in future_inputs] + future_embeddings_list = [embeds for _, embeds in future_inputs] layer_results = cast( MoEDecoderLayerMicroBatchOutput, @@ -286,6 +291,39 @@ def _micro_batch_forward( return outputs_per_mb + def _gather_future_embeddings(self, seq_ctx: SequenceContext) -> torch.Tensor | None: + if self.mtp_config.use_local_future_embeddings and seq_ctx.inputs_embeds is not None: + # Preserve the existing gather semantics. Access once per microbatch; + # the shared context retains one full base, never seven rolled bases. + return seq_ctx.raw_inputs_embeds + return None + + def _next_future_inputs( + self, + seq_ctx: SequenceContext, + embed_tokens_fn: Callable[[torch.Tensor], torch.Tensor], + *, + depth: int, + raw_inputs_embeds: torch.Tensor | None, + ) -> tuple[SequenceContext, torch.Tensor]: + if raw_inputs_embeds is None: + rolled_ctx = roll_sequence_context(seq_ctx, shifts=-1) + return rolled_ctx, self._embed_future(rolled_ctx, embed_tokens_fn) + + assert seq_ctx.inputs_embeds is not None + mesh = seq_ctx.sequence_parallel_mesh + shard_start = seq_ctx._shard_start if mesh is not None and mesh.size() > 1 else 0 + future = local_packed_future_embeddings( + raw_inputs_embeds, + seq_ctx.cu_seq_lens_q, + depth=depth, + shard_start=shard_start, + shard_size=seq_ctx.inputs_embeds.shape[1], + ) + # Decoder metadata (pack boundaries, positions, routing) is depth-invariant. + # The shifted embedding is an explicit input, so the context stays shared. + return seq_ctx, future + @staticmethod def _embed_future( seq_ctx: SequenceContext, diff --git a/xtuner/v1/module/mtp/utils.py b/xtuner/v1/module/mtp/utils.py index 3ed86e3831..75f057ce5b 100644 --- a/xtuner/v1/module/mtp/utils.py +++ b/xtuner/v1/module/mtp/utils.py @@ -5,6 +5,43 @@ from xtuner.v1.data_proto import SequenceContext +def local_packed_future_embeddings( + raw_inputs_embeds: torch.Tensor, + cu_seq_lens: torch.IntTensor, + *, + depth: int, + shard_start: int, + shard_size: int, +) -> torch.Tensor: + """Materialize an owned local embedding buffer for one MTP depth. + + For global query position ``i``, read ``i + depth`` only within the same + packed sequence, otherwise fill zero. This equals repeated pack-aware + left rolls followed by an SP slice, without allocating full rolled copies. + The full input uses the existing SequenceContext gather and its existing + autograd semantics; this function does not detach or communicate tensors. + """ + if depth < 1: + raise ValueError("MTP future depth must be positive") + if shard_start < 0 or shard_size < 0 or shard_start + shard_size > raw_inputs_embeds.shape[1]: + raise ValueError("Local MTP shard must lie within the full embedding sequence") + if shard_size == 0: + return raw_inputs_embeds[:, :0].clone() + + positions = torch.arange(shard_start, shard_start + shard_size, device=raw_inputs_embeds.device) + # right=True skips empty packs and assigns a boundary position to its next pack. + ends = cu_seq_lens[1:].to(device=positions.device) + pack_indices = torch.searchsorted(ends, positions, right=True) + pack_ends = ends.index_select(0, pack_indices) + source_positions = positions + depth + valid = source_positions < pack_ends + # Invalid positions still need an in-bounds index before masking (last SP rank). + source_indices = source_positions.clamp_max(raw_inputs_embeds.shape[1] - 1) + future = raw_inputs_embeds.index_select(1, source_indices) + future.masked_fill_(~valid[None, :, None], 0) + return future + + def roll_packed_tensor( tensor: torch.Tensor, cu_seq_lens: torch.IntTensor,