fix: Fixed the .item() crash on multi-element MTP loss tensors via a

.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"

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()

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)

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__]))