Fix MindSpeed NPU LoRA projections and Qwen3.8 PLE host gather - #209
Draft
addsubmuldiv wants to merge 6 commits into
Draft
addsubmuldiv wants to merge 6 commits into
addsubmuldiv wants to merge 6 commits into
Conversation
_gather_ple_rows_from_pinned dereferences the host pointer of the CPU-pinned n-gram table inside the device kernel, which relies on CUDA unified virtual addressing. On Ascend NPU the host pointer is treated as a device DDR address and the first forward aborts with CANN error 0x800000 (MTE DDR out of range) on all vector cores of the first PP stage. Gate the fast path to ids.device.type == 'cuda' so NPU takes the numerically equivalent plain-torch fallback. CUDA behavior is unchanged. Verified: 300-step NPU Megatron LoRA run of Qwen3.8-Flash-Next (TP2/EP4/ETP1/PP2, 8x910B3), per-step loss MAE vs GPU (8xH20) 0.00687.
…allel linears
Under bf16 (non-MC2) MindSpeed rebinds megatron.core.extensions.
transformer_engine.TEColumnParallelLinear/TERowParallelLinear to the
native mcore classes. The import names still resolve on both platforms,
so on NPU three assumptions of the CUDA-oriented LoRA path break:
- dispatch: isinstance against the TE tuple misses the native classes
and silently falls through to peft's vanilla branch ("Target module
ColumnParallelLinear not supported"). Enumerate the native classes
in the dispatch tuple; on CUDA they never match.
- construction: LoRA factors sized with TE shard semantics are built
by the native classes with global semantics, surfacing at first
forward as mismatches of exactly tp_size (aclnnMatmul EZ1001
k-axis [.,.,3072]x[6144,.]). Use the per-rank size on NPU.
- forward: the MindSpeed TELayerNormColumnParallelLinear returns
(result, bias) instead of ((result, layernorm_out), bias);
reconstruct the norm output via the base layer's _rmsnorm on NPU.
Complements 4450669 (NPU grouped-linear path). Known gap left open:
replicated TELinear bases (hyper-connection projections) still build
lora_b through a sharded column path; only reachable with
all-linear targets.
Verified: 300-step NPU Megatron LoRA run of Qwen3.8-Flash-Next
(TP2/EP4/ETP1/PP2, 8x910B3), per-step loss MAE vs GPU (8xH20) 0.00687.
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
Native row dispatch and classification remain incompatible across backend and alias paths.
Review effort: Lite
Findings: None
What changed in this PR
Enables MindSpeed NPU LoRA support for Qwen3.8 projections and avoids unsupported PLE Triton gathers on non-CUDA devices.
Changes:
- Dispatches native parallel linear layers to LoRA.
- Adjusts NPU row-projection sizing and forward handling.
- Uses the torch PLE fallback on non-CUDA devices.
| File | Summary |
|---|---|
src/mcore_bridge/tuners/patcher.py |
Adds native parallel dispatch; backend and row-classification compatibility issues remain. |
src/mcore_bridge/tuners/lora.py |
Supports NPU sizing and tuple outputs; native row projections need correct classification. |
src/mcore_bridge/model/modules/kernels/ple_kernels.py |
Restricts pinned-host Triton gathering to CUDA devices. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
addsubmuldiv
marked this pull request as draft
September 24, 2026 08:32
addsubmuldiv
marked this pull request as ready for review
September 24, 2026 09:20
addsubmuldiv
marked this pull request as draft
September 24, 2026 13:01
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.
Scope
Enable Qwen3.8-Flash-Next Megatron LoRA training on the MindSpeed NPU stack, including Swift's
all-lineartarget selection.ColumnParallelLinearandRowParallelLinearto bridge LoRA. Construct native Row LoRA A with the base layer's global input width andinput_is_parallelsetting, and handle its(output, bias)forward contract.TELinearbase on NPU, construct a replicated LoRA B with the full output width. MindSpeed'sTEColumnParallelLinearwould shard that width by TP. Mark both replicated factors for TP gradient synchronization under sequence parallelism.This PR does not change model definitions or the Swift Megatron training entrypoint.
Validation
On 8×Ascend 910B3, CANN 9.0, Megatron Core 0.16, MindSpeed 0.16, BF16, TP2/PP2/EP4/ETP1, PLE CPU offload, and the full local Qwen3.8-Flash-Next checkpoint:
mainin_proj out_proj linear_proj linear_qkvmainmainTarget module ColumnParallelLinear ... is not supported.mainall-linearinput_mix_weight_down: base output 320 vs LoRA output 160 with TP2.mainall-linearinput_mix_weight_down.lora_B.weightshape[320, 8]and nonzero values.The paired target runs used the original 300-step parity script with
train_iters=2: same model, text dataset, seed, max length 256, LoRA rank 8, micro batch 1, global batch 8, and all other training settings. Their recordedargs.jsonfiles differ only intarget_modulesamong the checked training settings. Both runs loaded weights, injected LoRA, and completed forward, backward, and optimizer steps.A separate TP2 NPU smoke directly targets native
RowParallelLinearthroughdispatch_megatron. Bothinput_is_parallel=FalseandTruepassed injection, forward, backward, an SGD step, and merge/unmerge; LoRA A, B, and merged delta had the expected local shapes(2, 4),(8, 2), and(8, 4).The full HF merged-weight export was stopped after training because this smoke only needed to validate training. Full export, resume, image/video training, and longer runs are not claimed. Python compile, flake8, YAPF formatting,
git diff --check, and PR lint passed.The open Core 0.18/TENPU migration PR #189 targets a different backend stack and is not validation of this MindSpeed 0.16 path.