Skip to content

[NPU] align fused linear JSD forward with accum_dtype API - #1346

Open
zheliuyu wants to merge 1 commit into
linkedin:mainfrom
zheliuyu:fix-npu-fused_linear_jsd_forward
Open

[NPU] align fused linear JSD forward with accum_dtype API#1346
zheliuyu wants to merge 1 commit into
linkedin:mainfrom
zheliuyu:fix-npu-fused_linear_jsd_forward

Conversation

@zheliuyu

@zheliuyu zheliuyu commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Summary

Upstream #1322 added accum_dtype to fused linear JSD (LigerFusedLinearJSD, LigerFusedLinearJSDFunction.apply). On NPU, liger_kernel.ops resolves to the Ascend backend, whose LigerFusedLinearJSDFunction.forward still used the old signature. All test/transformers/test_fused_linear_jsd.py cases failed with:

TypeError: forward() takes from 5 to 9 positional arguments but 10 were given

This PR updates backends/_ascend/ops/fused_linear_jsd.py to accept accum_dtype, pass it through fused_linear_jsd_forward, and mirror the default path: optional fp32 grad_weight accumulation with cast back to weight dtype. backward return tuple is extended by one None for the new argument.

Testing Done

image
  • Hardware Type: All NPUs
  • run make test to ensure correctness
  • run make checkstyle to ensure code style
  • run make test-convergence to ensure convergence

@zheliuyu

zheliuyu commented Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

@Tcc0403 This PR is ready for review. Thanks.

@Tcc0403 Tcc0403 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM

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.

2 participants