From 3619ce90cc6e8692ef5d6af2b95bef4f39386a48 Mon Sep 17 00:00:00 2001 From: Amitesh Gupta Date: Mon, 6 Jul 2026 12:38:56 +0530 Subject: [PATCH] fix: Fixed the `.item()` crash on multi-element MTP loss tensors via a Addresses #2131. --- .github/workflows/pr-test.yml | 4 ++ slime/backends/megatron_utils/cp_utils.py | 10 +++++ slime/backends/megatron_utils/model.py | 16 +++++--- tests/test_mtp_loss_logging.py | 48 +++++++++++++++++++++++ 4 files changed, 73 insertions(+), 5 deletions(-) create mode 100644 tests/test_mtp_loss_logging.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 761496b949..292d4951a8 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -573,6 +573,10 @@ jobs: "num_gpus": 0, "test_file": "test_cp_utils.py" }, + { + "num_gpus": 0, + "test_file": "test_mtp_loss_logging.py" + }, { "num_gpus": 0, "test_file": "test_metric_report.py" diff --git a/slime/backends/megatron_utils/cp_utils.py b/slime/backends/megatron_utils/cp_utils.py index 96c97df0e2..6b0531a6b7 100644 --- a/slime/backends/megatron_utils/cp_utils.py +++ b/slime/backends/megatron_utils/cp_utils.py @@ -403,3 +403,13 @@ def prepare_routed_experts_for_routing_replay( routed_experts = routed_experts[start:end] return routed_experts + + +def compute_mtp_losses(values: torch.Tensor, mtp_loss_scale: float) -> list[float]: + """Scale the MTP loss tracker's ``values`` into one loss per MTP layer. + + ``values`` holds one accumulated loss per MTP layer (shape + ``(mtp_num_layers,)``), so a multi-head MTP model (``mtp_num_layers > 1``) + must not be squeezed into a single scalar via ``.item()``. + """ + return (values * mtp_loss_scale).tolist() diff --git a/slime/backends/megatron_utils/model.py b/slime/backends/megatron_utils/model.py index 1ad6cd7957..91572bb7a3 100644 --- a/slime/backends/megatron_utils/model.py +++ b/slime/backends/megatron_utils/model.py @@ -32,7 +32,7 @@ from slime.utils.memory_utils import clear_memory from .checkpoint import load_checkpoint, save_checkpoint -from .cp_utils import reduce_train_step_metrics +from .cp_utils import compute_mtp_losses, reduce_train_step_metrics from .data import DataIterator, get_batch from .loss import ROLLOUT_TOP_P_TOKEN_KEYS, get_rollout_top_p_logprob_kwargs, loss_function from .model_provider import get_model_provider_func @@ -854,15 +854,16 @@ def train( torch.distributed.all_reduce(values, group=tracker.get("reduce_group")) if tracker.get("avg_group") is not None: torch.distributed.all_reduce(values, group=tracker["avg_group"], op=torch.distributed.ReduceOp.AVG) - # here we assume only one mtp layer - mtp_losses = (tracker["values"] * mtp_loss_scale).item() + # one loss value per MTP layer (values has mtp_num_layers elements) + mtp_losses = compute_mtp_losses(tracker["values"], mtp_loss_scale) MTPLossLoggingHelper.clean_loss_in_tracker() # CI check: verify MTP loss is within expected bounds if args.ci_test: from slime.backends.megatron_utils.ci_utils import check_mtp_loss - check_mtp_loss(mtp_losses) + for mtp_loss in mtp_losses: + check_mtp_loss(mtp_loss) # per train step log. if ( @@ -879,7 +880,12 @@ def train( } log_dict[f"train/{role_tag}grad_norm"] = grad_norm if args.enable_mtp_training: - log_dict[f"train/{role_tag}mtp_loss"] = mtp_losses + # keep the pre-existing single aggregate key for dashboards/alerts that + # already depend on it, and add a per-layer breakdown for multi-head MTP. + log_dict[f"train/{role_tag}mtp_loss"] = sum(mtp_losses) / len(mtp_losses) + if len(mtp_losses) > 1: + for layer_idx, mtp_loss in enumerate(mtp_losses): + log_dict[f"train/{role_tag}mtp_loss_{layer_idx}"] = mtp_loss for param_group_id, param_group in enumerate(optimizer.param_groups): log_dict[f"train/{role_tag}lr-pg_{param_group_id}"] = opt_param_scheduler.get_lr(param_group) diff --git a/tests/test_mtp_loss_logging.py b/tests/test_mtp_loss_logging.py new file mode 100644 index 0000000000..0e1bf4c076 --- /dev/null +++ b/tests/test_mtp_loss_logging.py @@ -0,0 +1,48 @@ +"""CPU unit tests for ``slime.backends.megatron_utils.cp_utils.compute_mtp_losses``. + +Regression test for https://github.com/THUDM/slime/issues/2131: multi-head +MTP models (``--mtp-num-layers > 1``) crashed at per-step logging because +the MTP-loss tracker's ``values`` tensor (one entry per MTP layer) was +squeezed through ``.item()``, which only works for a single-element tensor. + +The CPU-only CI image does not ship megatron -- ``_cp_dist_helpers`` stubs +``megatron.core.mpu`` at import time so the subsequent ``cp_utils`` import +binds against the stub. +""" + +from __future__ import annotations + +# Import the helpers BEFORE the slime imports so the megatron stub lands +# in sys.modules first. pytest's prepend importmode puts this file's +# directory (``tests/``) on sys.path, which is what makes the bare-name +# import work without an ``__init__.py``. +import _cp_dist_helpers # noqa: F401 +import pytest +import torch + +from slime.backends.megatron_utils.cp_utils import compute_mtp_losses # noqa: E402 + + +NUM_GPUS = 0 + + +@pytest.mark.unit +def test_single_mtp_layer_returns_one_loss(): + """``mtp_num_layers == 1`` keeps producing a single scaled loss value.""" + values = torch.tensor([2.0]) + assert compute_mtp_losses(values, mtp_loss_scale=0.5) == pytest.approx([1.0]) + + +@pytest.mark.unit +def test_multi_mtp_layer_does_not_crash(): + """``mtp_num_layers > 1`` must not raise -- this is the reported crash: + calling ``.item()`` on a multi-element tensor raises + ``RuntimeError: a Tensor with 3 elements cannot be converted to Scalar``. + """ + values = torch.tensor([2.0, 4.0, 6.0]) + losses = compute_mtp_losses(values, mtp_loss_scale=0.5) + assert losses == pytest.approx([1.0, 2.0, 3.0]) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__]))