Skip to content

[Fix] Fix device mismatch for ascend npu when loss_cfg.loss_reduction != "token" #1365

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

Jump to bottom
Open
nil0x9 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[Fix] Fix device mismatch for ascend npu when loss_cfg.loss_reduction != "token" #1365

nil0x9 wants to merge 1 commit into from

Conversation

@nil0x9
Copy link
Contributor

@nil0x9 nil0x9 commented Dec 15, 2025 •
edited
Loading

Currently when one passes loss_cfg.loss_reduction other than "token" on a ascend/npu device, a Runtime Error (device mismatch) is expected in this line:

xtuner/xtuner/v1/loss/ce_loss.py

Line 191 in d339718

loss = (loss * loss_weight).sum()

The root cause of this error is that, in ascend npu device, cu_seq_lens tensors are required to be on cpu. In func build_batches_loss_kwargs, the devuce ofloss_weight is inherited from num_grad_tokens -> boundaries -> cu_seq_lens -- and hence the problem.
@nil0x9 nil0x9 marked this pull request as ready for review December 15, 2025 17:38
@nil0x9 nil0x9 force-pushed the linty/fix-npu-loss-weight-device-mismatch branch from 8613e2d to f1294d3 Compare December 16, 2025 15:40
@nil0x9 nil0x9 force-pushed the linty/fix-npu-loss-weight-device-mismatch branch from f1294d3 to dacb7ef Compare December 24, 2025 08:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

2 participants

@nil0x9 @HAOCHENYE