Skip to content

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

Open
singlaamitesh wants to merge 1 commit into
THUDM:mainfrom
singlaamitesh:fix/2131-slime-2131
Open

fix: Fixed the .item() crash on multi-element MTP loss tensors via a#2182
singlaamitesh wants to merge 1 commit into
THUDM:mainfrom
singlaamitesh:fix/2131-slime-2131

Conversation

@singlaamitesh

Copy link
Copy Markdown

Fixes #2131

Summary

Fixed the .item() crash on multi-element MTP loss tensors via a new cp_utils.compute_mtp_losses helper (per-layer list) used for the CI bounds check and logging; after self-review, also fixed a follow-on regression where the multi-layer path had dropped the aggregate train/mtp_loss key that existing dashboards/alerts depend on — it now always logs the mean as mtp_loss plus per-layer mtp_loss_{i} keys when there's more than one MTP head.

Changes

  • slime/backends/megatron_utils/model.py
  • slime/backends/megatron_utils/cp_utils.py
  • tests/test_mtp_loss_logging.py
  • .github/workflows/pr-test.yml

How this was tested

Ran go=github)](https://github.com/thudm/slime/pull/2053/checks) locally.

Copilot AI review requested due to automatic review settings July 6, 2026 07:17

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] Multi-head MTP (--mtp-num-layers > 1) crashes at training-step logging

2 participants