[Feature] Support Intern-S2-Preview training on NPU - #2052
Open
CyCle1024 wants to merge 7 commits into
Open
Conversation
CyCle1024
force-pushed
the
ccy/npu/intern_s2_preview_support
branch
2 times, most recently
from
August 28, 2026 06:59
f30f490 to
95043dc
Compare
CyCle1024
commented
Aug 28, 2026
CyCle1024
commented
Aug 28, 2026
CyCle1024
commented
Aug 28, 2026
CyCle1024
commented
Aug 28, 2026
CyCle1024
commented
Aug 28, 2026
CyCle1024
commented
Aug 28, 2026
CyCle1024
commented
Aug 28, 2026
CyCle1024
commented
Aug 28, 2026
| layer_hidden_states = hidden_states | ||
| hidden_states = self.norm(hidden_states) | ||
|
|
||
| # Get LM loss context from dict |
CyCle1024
commented
Aug 28, 2026
CyCle1024
commented
Aug 28, 2026
| # leave the first post-MTP LM-head call to unshard on demand for now. A future | ||
| # optimization can explicitly prefetch the LM head in MTPBlock only after its final | ||
| # logical depth finishes. | ||
| if not mtp_config.share_weights: |
CyCle1024
commented
Aug 28, 2026
CyCle1024
commented
Aug 28, 2026
| nonpad_indices = seq_ctx.nonpad_indices | ||
| non_pad_token = nonpad_indices.numel() | ||
| num_tokens_global, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device) | ||
| num_tokens_global, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, nonpad_indices.device) |
CyCle1024
commented
Aug 28, 2026
CyCle1024
force-pushed
the
ccy/npu/intern_s2_preview_support
branch
from
August 30, 2026 09:59
83a2861 to
c64718d
Compare
Co-authored-by: CyCle1024 <cycle1024@gmail.com>
Co-authored-by: houyufeng4@huawei.com <houyufeng4@huawei.com> Co-authored-by: CyCle1024 <cycle1024@gmail.com>
Co-authored-by: CyCle1024 <cycle1024@gmail.com>
Co-authored-by: CyCle1024 <cycle1024@gmail.com>
Co-authored-by: CyCle1024 <cycle1024@gmail.com>
Co-authored-by: CyCle1024 <cycle1024@gmail.com>
CyCle1024
force-pushed
the
ccy/npu/intern_s2_preview_support
branch
from
August 30, 2026 10:53
c64718d to
a0381c6
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds NPU training support for Intern-S2-Preview while keeping the shared model and CUDA operator contracts stable.
Design notes
[B, T, H, D]as its public layout. The NPU wrappers preserve native head-first storage with transpose views where possible and materialize only when required by a fixed-layout kernel.xtuner/v1/ops; model modules keep device-independent logical interfaces.Known limitations
skip_dispatch_pad_tokenscurrently covers the standard single-context MoE forward, not the intra-layer multi-microbatch/domino-EP path.B == 1,num_value_headsdivisible bynum_key_heads, and equal key/value head dimensions.export_bind_config(). Trainerrun()uses the packaged binding configuration; host IRQ setup remains an explicit launch-script operation.