Skip to content

[Feature] Support Intern-S2-Preview training on NPU - #2052

Open
CyCle1024 wants to merge 7 commits into
InternLM:mainfrom
CyCle1024:ccy/npu/intern_s2_preview_support
Open

[Feature] Support Intern-S2-Preview training on NPU#2052
CyCle1024 wants to merge 7 commits into
InternLM:mainfrom
CyCle1024:ccy/npu/intern_s2_preview_support

Conversation

@CyCle1024

@CyCle1024 CyCle1024 commented Aug 27, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR adds NPU training support for Intern-S2-Preview while keeping the shared model and CUDA operator contracts stable.

  • Improve MTP recomputation, FSDP scheduling, parameter residency, and LM-head prefetch control.
  • Preserve device Tensor sequence boundaries and pass NPU CPU-list metadata through explicit parameters.
  • Add Triton-Ascend fused cross-entropy for large-vocabulary training.
  • Skip tail-padding tokens in the standard MoE dispatch path and cache the required indices.
  • Add self-contained NPU GatedDeltaNet causal-convolution, gated-delta-rule, RMSNormGated, metadata, and sequence-parallel support.
  • Integrate NPU CPU-affinity binding into the trainer lifecycle and retain host IRQ configuration helpers.
  • Remove the dependency on a version-specific private DTensor shard API.

Design notes

  • Public cumulative-length arguments remain device tensors. NPU-only CPU lists and chunk metadata use separate, explicit parameters, so CUDA kernels do not receive Python lists.
  • GatedDeltaNet uses [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.
  • Normal GatedDeltaNet forward keeps one mixed-QKV convolution. Sequence-parallel forward gathers the global packed-token axis, shards heads, and runs separate Q/K/V convolutions, matching the NPU kernel geometry.
  • GatedDeltaNet metadata is prepared once per sequence context and reused across layers. Cache reuse requires identical cumulative sequence boundaries, device, token count, and kernel block configuration.
  • Backend-specific implementations and dispatch stay under xtuner/v1/ops; model modules keep device-independent logical interfaces.
  • CUDA keeps its existing tensor inputs, loss-reduction behavior, and time-major GatedDeltaNet contract.

Known limitations

  • skip_dispatch_pad_tokens currently covers the standard single-context MoE forward, not the intra-layer multi-microbatch/domino-EP path.
  • GatedDeltaNet currently requires B == 1, num_value_heads divisible by num_key_heads, and equal key/value head dimensions.
  • In normal non-SP NPU GatedDeltaNet backward, PyTorch's generic mixed-QKV split backward produces a time-major contiguous gradient. The fixed-stride causal-convolution backward kernel therefore needs one additional layout materialization compared with the native-layout Tina prototype. The SP path does not have this mixed-QKV split. A stride-preserving custom autograd split is the intended follow-up optimization.
  • Dynamic device-task and IRQ discovery is used by export_bind_config(). Trainer run() uses the packaged binding configuration; host IRQ setup remains an explicit launch-script operation.

@CyCle1024 CyCle1024 changed the title [Feature] Add comprehensive NPU support for GatedDeltaNet [WIP][Feature] Add comprehensive NPU support for GatedDeltaNet Aug 28, 2026
@CyCle1024
CyCle1024 force-pushed the ccy/npu/intern_s2_preview_support branch 2 times, most recently from f30f490 to 95043dc Compare August 28, 2026 06:59
Comment thread xtuner/v1/train/trainer.py Outdated
Comment thread xtuner/v1/optim/muon.py Outdated
Comment thread xtuner/v1/ops/rms_norm/__init__.py Outdated
Comment thread xtuner/v1/data_proto/sequence_context.py Outdated
Comment thread xtuner/v1/data_proto/sequence_context.py Outdated
Comment thread xtuner/v1/model/moe/moe.py
Comment thread xtuner/v1/model/moe/moe.py Outdated
layer_hidden_states = hidden_states
hidden_states = self.norm(hidden_states)

# Get LM loss context from dict

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

rollback

Comment thread xtuner/v1/model/moe/moe.py
Comment thread xtuner/v1/model/moe/moe.py Outdated
# 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:

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

rollback

Comment thread xtuner/v1/model/moe/moe.py
Comment thread xtuner/v1/model/moe/qwen3vl_text.py Outdated
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)

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

rollback

Comment thread xtuner/v1/module/dispatcher/torch_all2all.py Outdated

@CyCle1024 CyCle1024 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

fix

@CyCle1024 CyCle1024 changed the title [WIP][Feature] Add comprehensive NPU support for GatedDeltaNet [WIP][Feature] Add InternS2_35B_preview NPU support, along with some general optimization Aug 28, 2026
@CyCle1024
CyCle1024 force-pushed the ccy/npu/intern_s2_preview_support branch from 83a2861 to c64718d Compare August 30, 2026 09:59
Feng0w0 and others added 7 commits August 30, 2026 18:45
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
CyCle1024 requested a review from jayhenry August 30, 2026 10:51
@CyCle1024
CyCle1024 force-pushed the ccy/npu/intern_s2_preview_support branch from c64718d to a0381c6 Compare August 30, 2026 10:53
@CyCle1024 CyCle1024 changed the title [WIP][Feature] Add InternS2_35B_preview NPU support, along with some general optimization [Feature] Add InternS2_35B_preview NPU support, along with some general optimization Aug 30, 2026
@CyCle1024 CyCle1024 changed the title [Feature] Add InternS2_35B_preview NPU support, along with some general optimization [Feature] Support Intern-S2-Preview training on NPU Aug 30, 2026
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