From 398479a913fa08720960d786b6122ecbab22cfa3 Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Thu, 24 Sep 2026 06:47:25 +0000 Subject: [PATCH 1/6] fix: restrict PLE pinned-host triton fast path to CUDA (NPU lacks UVA) _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. --- .../model/modules/kernels/ple_kernels.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/mcore_bridge/model/modules/kernels/ple_kernels.py b/src/mcore_bridge/model/modules/kernels/ple_kernels.py index 28fb9ac..979c36b 100644 --- a/src/mcore_bridge/model/modules/kernels/ple_kernels.py +++ b/src/mcore_bridge/model/modules/kernels/ple_kernels.py @@ -60,12 +60,15 @@ def gather_ple_rows(host_table, ids, row_start, row_end, out=None): row_start / row_end: this rank's global row range. out: optional preallocated ``[*ids.shape, embedding_dim]`` bf16 device tensor. """ - # Gate on "triton exists + ids live on an accelerator", not on is_cuda: - # triton ships per-vendor backends (ROCm in-tree; NPU/XPU via vendor forks - # such as triton-ascend), so a platform that cannot run this kernel fails - # at launch instead of being silently excluded here. The device-side read - # of the pinned host table is only verified on CUDA so far. - if not HAVE_TRITON or ids.device.type == 'cpu': + # Gate on "triton exists + ids live on CUDA/ROCm". The device-side read of + # the pinned host table relies on CUDA unified addressing: the kernel + # dereferences the host-pinned data_ptr directly. triton-ascend on NPU + # happily launches the same kernel, but the host pointer is not a + # device-mapped DDR address there and the load faults (aivec "DDR address + # of the MTE instruction is out of range", all vector cores at the same + # pc). Fall back to the numerically identical torch path on any other + # device type. + if not HAVE_TRITON or ids.device.type != 'cuda': return None if host_table.dtype != torch.bfloat16: return None From 27bb52b59be2bc7704c0bb9b9d3d3da2655b1c17 Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Thu, 24 Sep 2026 06:47:32 +0000 Subject: [PATCH 2/6] fix: LoRA dispatch and construction for MindSpeed-replaced native parallel 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. --- src/mcore_bridge/tuners/lora.py | 17 ++++++++++++++--- src/mcore_bridge/tuners/patcher.py | 8 +++++++- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/src/mcore_bridge/tuners/lora.py b/src/mcore_bridge/tuners/lora.py index 7582650..0f0ebe0 100644 --- a/src/mcore_bridge/tuners/lora.py +++ b/src/mcore_bridge/tuners/lora.py @@ -16,6 +16,7 @@ TERowParallelGroupedLinear, TERowParallelLinear) from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding from megatron.core.parallel_state import get_expert_tensor_parallel_world_size, get_tensor_model_parallel_world_size +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear from megatron.core.tensor_parallel.random import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name from megatron.core.transformer.mlp import apply_swiglu_sharded_factory from megatron.core.transformer.module import MegatronModule @@ -215,15 +216,19 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): **kwargs, ) else: + # MindSpeed aliases TERowParallelLinear to the native mcore class, + # whose input_size is global (sharded internally), unlike TE where + # the per-shard size is passed. Reuse the full size on NPU. + row_input_size = self.in_features if is_torch_npu_available() else in_features lora_a = TERowParallelLinear( - input_size=in_features, + input_size=row_input_size, output_size=r, bias=False, input_is_parallel=True, **kwargs, ) lora_b = _build_local_te_linear(r, self.out_features, lora_bias, **kwargs) - lora_a.parallel_mode = self.base_layer.parallel_mode # fix moe_shared_expert_overlap + lora_a.parallel_mode = getattr(self.base_layer, 'parallel_mode', None) # fix moe_shared_expert_overlap else: if is_torch_npu_available(): out_features = self.out_features @@ -269,7 +274,7 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): gather_output=False, **kwargs, ) - lora_b.parallel_mode = self.base_layer.parallel_mode # fix moe_shared_expert_overlap + lora_b.parallel_mode = getattr(self.base_layer, 'parallel_mode', None) # fix moe_shared_expert_overlap for lora in [lora_a, lora_b]: # When parallel_mode is set to None by moe_shared_expert_overlap, # disable UB comm overlap; the corresponding collectives are driven @@ -427,6 +432,12 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any): (result, x), bias = self.base_layer(x, *args, **kwargs) elif isinstance(self.base_layer, (TELinear, TEGroupedLinear)): result, bias = self.base_layer(x, *args, **kwargs) + elif isinstance(self.base_layer, (ColumnParallelLinear, RowParallelLinear)): + # Native mcore parallel linears: MindSpeed on NPU aliases the TE + # column/row classes to these, so a spec written with + # TEColumnParallelLinear/TERowParallelLinear produces native layers + # there. Their forward also returns (output, output_bias). + result, bias = self.base_layer(x, *args, **kwargs) elif isinstance(self.base_layer, TopKRouter): with self._patch_router_gating(): result, bias = self.base_layer(x, *args, **kwargs) diff --git a/src/mcore_bridge/tuners/patcher.py b/src/mcore_bridge/tuners/patcher.py index 4a7f8e0..a34bbde 100644 --- a/src/mcore_bridge/tuners/patcher.py +++ b/src/mcore_bridge/tuners/patcher.py @@ -1,5 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from megatron.core.extensions.transformer_engine import TEGroupedLinear, TELayerNormColumnParallelLinear, TELinear +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.moe.router import TopKRouter from peft import LoraModel @@ -27,7 +28,12 @@ def dispatch_megatron( else: target_base_layer = target - linear_cls = (TELayerNormColumnParallelLinear, TELinear, TEGroupedLinear, TopKRouter) + # MindSpeed aliases TEColumnParallelLinear/TERowParallelLinear to the native + # mcore classes (breaking the TELinear inheritance used on CUDA), so the + # native classes must be dispatched as well. On CUDA these never match: + # every LoRA-capable projection there is built as a TE class. + linear_cls = (TELayerNormColumnParallelLinear, TELinear, TEGroupedLinear, TopKRouter, ColumnParallelLinear, + RowParallelLinear) if isinstance(target_base_layer, linear_cls): new_module = LoraParallelLinear(base_layer=target, adapter_name=adapter_name, **kwargs) From e5038eaa4feb1d941a003951c92025de0a147e82 Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Thu, 24 Sep 2026 09:03:01 +0000 Subject: [PATCH 3/6] fix: preserve replicated NPU LoRA output across tensor parallel ranks --- src/mcore_bridge/tuners/lora.py | 31 ++++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/src/mcore_bridge/tuners/lora.py b/src/mcore_bridge/tuners/lora.py index 0f0ebe0..46420ee 100644 --- a/src/mcore_bridge/tuners/lora.py +++ b/src/mcore_bridge/tuners/lora.py @@ -166,6 +166,8 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): self.lora_dropout[adapter_name] = lora_dropout_layer + replicated_base = (is_torch_npu_available() and isinstance(self.base_layer, TELinear) + and getattr(self.base_layer, 'parallel_mode', None) == 'duplicated') # lora needs to be forced to upgrade to 32-bit precision, otherwise it will overflow kwargs = { 'skip_bias_add': False, @@ -267,14 +269,19 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): ) else: lora_a = _build_local_te_linear(self.in_features, r, lora_bias, **kwargs) - lora_b = TEColumnParallelLinear( - input_size=r, - output_size=out_features, - bias=lora_bias, - gather_output=False, - **kwargs, - ) - lora_b.parallel_mode = getattr(self.base_layer, 'parallel_mode', None) # fix moe_shared_expert_overlap + if replicated_base: + # MindSpeed's column linear splits the output by TP, while + # a duplicated TELinear keeps the full output on each rank. + lora_b = _build_local_te_linear(r, self.out_features, lora_bias, **kwargs) + else: + lora_b = TEColumnParallelLinear( + input_size=r, + output_size=out_features, + bias=lora_bias, + gather_output=False, + **kwargs, + ) + lora_b.parallel_mode = getattr(self.base_layer, 'parallel_mode', None) # fix moe_shared_expert_overlap for lora in [lora_a, lora_b]: # When parallel_mode is set to None by moe_shared_expert_overlap, # disable UB comm overlap; the corresponding collectives are driven @@ -292,11 +299,13 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): # finalize_model_grads for parameters flagged `sequence_parallel` (same as layernorm # weights); without the flag each TP rank trains a different copy and export_weights # saves rank 0 only (observed: last layer linear_proj.lora_B saved as all zeros). + # For a duplicated base, both factors see the local sequence shard. if (self.tp_size > 1 and not isinstance(self.base_layer, TopKRouter) and (getattr(self.config, 'sequence_parallel', False) or self.sequence_parallel)): - replicated = lora_b if self.is_parallel_a else lora_a - for p in replicated.parameters(): - p.sequence_parallel = True + replicated_factors = (lora_a, lora_b) if replicated_base else (lora_b if self.is_parallel_a else lora_a, ) + for factor in replicated_factors: + for p in factor.parameters(): + p.sequence_parallel = True self.lora_A[adapter_name] = lora_a self.lora_B[adapter_name] = lora_b if hasattr(self, 'lora_bias'): From 64ab83c37c4485920f80ab799dfe94ba74f7ea3e Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Thu, 24 Sep 2026 09:05:57 +0000 Subject: [PATCH 4/6] style: format replicated LoRA handling --- src/mcore_bridge/tuners/lora.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/mcore_bridge/tuners/lora.py b/src/mcore_bridge/tuners/lora.py index 46420ee..0507433 100644 --- a/src/mcore_bridge/tuners/lora.py +++ b/src/mcore_bridge/tuners/lora.py @@ -166,8 +166,9 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): self.lora_dropout[adapter_name] = lora_dropout_layer - replicated_base = (is_torch_npu_available() and isinstance(self.base_layer, TELinear) - and getattr(self.base_layer, 'parallel_mode', None) == 'duplicated') + replicated_base = ( + is_torch_npu_available() and isinstance(self.base_layer, TELinear) + and getattr(self.base_layer, 'parallel_mode', None) == 'duplicated') # lora needs to be forced to upgrade to 32-bit precision, otherwise it will overflow kwargs = { 'skip_bias_add': False, @@ -281,7 +282,8 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): gather_output=False, **kwargs, ) - lora_b.parallel_mode = getattr(self.base_layer, 'parallel_mode', None) # fix moe_shared_expert_overlap + lora_b.parallel_mode = getattr(self.base_layer, 'parallel_mode', + None) # fix moe_shared_expert_overlap for lora in [lora_a, lora_b]: # When parallel_mode is set to None by moe_shared_expert_overlap, # disable UB comm overlap; the corresponding collectives are driven From 478b1d61d9aaf65442a93acdc7c3b625f2e18b34 Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Thu, 24 Sep 2026 09:30:55 +0000 Subject: [PATCH 5/6] docs: shorten NPU LoRA compatibility comments --- .../model/modules/kernels/ple_kernels.py | 9 +------- src/mcore_bridge/tuners/lora.py | 22 +++++-------------- src/mcore_bridge/tuners/patcher.py | 5 +---- 3 files changed, 7 insertions(+), 29 deletions(-) diff --git a/src/mcore_bridge/model/modules/kernels/ple_kernels.py b/src/mcore_bridge/model/modules/kernels/ple_kernels.py index 979c36b..57381a3 100644 --- a/src/mcore_bridge/model/modules/kernels/ple_kernels.py +++ b/src/mcore_bridge/model/modules/kernels/ple_kernels.py @@ -60,14 +60,7 @@ def gather_ple_rows(host_table, ids, row_start, row_end, out=None): row_start / row_end: this rank's global row range. out: optional preallocated ``[*ids.shape, embedding_dim]`` bf16 device tensor. """ - # Gate on "triton exists + ids live on CUDA/ROCm". The device-side read of - # the pinned host table relies on CUDA unified addressing: the kernel - # dereferences the host-pinned data_ptr directly. triton-ascend on NPU - # happily launches the same kernel, but the host pointer is not a - # device-mapped DDR address there and the load faults (aivec "DDR address - # of the MTE instruction is out of range", all vector cores at the same - # pc). Fall back to the numerically identical torch path on any other - # device type. + # Pinned-host fast path is CUDA-only; use the torch fallback elsewhere. if not HAVE_TRITON or ids.device.type != 'cuda': return None if host_table.dtype != torch.bfloat16: diff --git a/src/mcore_bridge/tuners/lora.py b/src/mcore_bridge/tuners/lora.py index 0507433..df3d7e8 100644 --- a/src/mcore_bridge/tuners/lora.py +++ b/src/mcore_bridge/tuners/lora.py @@ -219,9 +219,7 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): **kwargs, ) else: - # MindSpeed aliases TERowParallelLinear to the native mcore class, - # whose input_size is global (sharded internally), unlike TE where - # the per-shard size is passed. Reuse the full size on NPU. + # Native NPU RowParallelLinear takes the global input size. row_input_size = self.in_features if is_torch_npu_available() else in_features lora_a = TERowParallelLinear( input_size=row_input_size, @@ -271,8 +269,7 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): else: lora_a = _build_local_te_linear(self.in_features, r, lora_bias, **kwargs) if replicated_base: - # MindSpeed's column linear splits the output by TP, while - # a duplicated TELinear keeps the full output on each rank. + # Match the base layer's replicated output. lora_b = _build_local_te_linear(r, self.out_features, lora_bias, **kwargs) else: lora_b = TEColumnParallelLinear( @@ -294,14 +291,8 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): lora.ub_overlap_ag_fprop = False lora.ub_overlap_rs_dgrad = False - # With sequence parallelism the replicated (non-sharded) LoRA factor only sees this TP - # rank's sequence shard: for RowParallel targets lora_A reduce-scatters its output before - # lora_B, and for ColumnParallel targets lora_A consumes the sequence-sharded input. Its - # gradient must therefore be summed over the TP group. Megatron does this in - # finalize_model_grads for parameters flagged `sequence_parallel` (same as layernorm - # weights); without the flag each TP rank trains a different copy and export_weights - # saves rank 0 only (observed: last layer linear_proj.lora_B saved as all zeros). - # For a duplicated base, both factors see the local sequence shard. + # Sequence-parallel inputs need TP reduction for replicated LoRA weights. + # Both factors are replicated when the base layer is duplicated. if (self.tp_size > 1 and not isinstance(self.base_layer, TopKRouter) and (getattr(self.config, 'sequence_parallel', False) or self.sequence_parallel)): replicated_factors = (lora_a, lora_b) if replicated_base else (lora_b if self.is_parallel_a else lora_a, ) @@ -444,10 +435,7 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any): elif isinstance(self.base_layer, (TELinear, TEGroupedLinear)): result, bias = self.base_layer(x, *args, **kwargs) elif isinstance(self.base_layer, (ColumnParallelLinear, RowParallelLinear)): - # Native mcore parallel linears: MindSpeed on NPU aliases the TE - # column/row classes to these, so a spec written with - # TEColumnParallelLinear/TERowParallelLinear produces native layers - # there. Their forward also returns (output, output_bias). + # Native parallel linears return (output, bias). result, bias = self.base_layer(x, *args, **kwargs) elif isinstance(self.base_layer, TopKRouter): with self._patch_router_gating(): diff --git a/src/mcore_bridge/tuners/patcher.py b/src/mcore_bridge/tuners/patcher.py index a34bbde..85a4aeb 100644 --- a/src/mcore_bridge/tuners/patcher.py +++ b/src/mcore_bridge/tuners/patcher.py @@ -28,10 +28,7 @@ def dispatch_megatron( else: target_base_layer = target - # MindSpeed aliases TEColumnParallelLinear/TERowParallelLinear to the native - # mcore classes (breaking the TELinear inheritance used on CUDA), so the - # native classes must be dispatched as well. On CUDA these never match: - # every LoRA-capable projection there is built as a TE class. + # MindSpeed uses native mcore classes for TE parallel linears. linear_cls = (TELayerNormColumnParallelLinear, TELinear, TEGroupedLinear, TopKRouter, ColumnParallelLinear, RowParallelLinear) if isinstance(target_base_layer, linear_cls): From 700e31a711bb3c4ff49fdfd7010434d00ee554e4 Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Thu, 24 Sep 2026 10:47:17 +0000 Subject: [PATCH 6/6] fix: handle native row-parallel LoRA targets --- src/mcore_bridge/tuners/lora.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/mcore_bridge/tuners/lora.py b/src/mcore_bridge/tuners/lora.py index df3d7e8..4412b25 100644 --- a/src/mcore_bridge/tuners/lora.py +++ b/src/mcore_bridge/tuners/lora.py @@ -118,7 +118,8 @@ def __init__( if use_dora: raise ValueError(f'{self.__class__.__name__} does not support DoRA yet, please set it to False') - self.is_parallel_a = isinstance(base_layer, (TERowParallelLinear, TERowParallelGroupedLinear)) + self.is_parallel_a = isinstance(base_layer, + (TERowParallelLinear, TERowParallelGroupedLinear, RowParallelLinear)) self.is_grouped = isinstance(base_layer, TEGroupedLinear) self.fan_in_fan_out = fan_in_fan_out self._active_adapter = adapter_name @@ -219,13 +220,14 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): **kwargs, ) else: - # Native NPU RowParallelLinear takes the global input size. - row_input_size = self.in_features if is_torch_npu_available() else in_features - lora_a = TERowParallelLinear( + native_row = isinstance(self.base_layer, RowParallelLinear) + row_input_size = self.in_features if native_row or is_torch_npu_available() else in_features + row_linear_cls = RowParallelLinear if native_row else TERowParallelLinear + lora_a = row_linear_cls( input_size=row_input_size, output_size=r, bias=False, - input_is_parallel=True, + input_is_parallel=getattr(self.base_layer, 'input_is_parallel', True), **kwargs, ) lora_b = _build_local_te_linear(r, self.out_features, lora_bias, **kwargs)