[Perf] Reduce communication overhead, improve NPU compatibility, and add CPU core binding

.dev_scripts/npu_bind_irq.sh

@@ -0,0 +1,89 @@
+
+
+all_npu_id=(0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15)
+
+all_start_cpus=(13 27 53 67 93 107 133 147 173 187 213 227 253 267 293 307)   
+
+# 将10进制CPU号转换为16进制mask码
+cpu_to_mask() {
+    local cpu=$1
+    local group=$((cpu / 32))
+    local bit=$((cpu % 32))
+    local value=$((1 << bit))
+    local mask=$(printf "%08x" $value)
+
+    # 拼接成逗号分隔的掩码字符串
+    for ((i=1; i<=group; i++)); do
+        mask="${mask},00000000"
+    done
+    echo $mask
+}
+
+check_and_stop_irqbalance() {
+    # 检查 irqbalance 是否存在
+    if systemctl list-unit-files | grep -q irqbalance.service; then
+        if systemctl is-active --quiet irqbalance; then
+            echo "检测到 irqbalance 服务正在运行，正在关闭..."
+            systemctl stop irqbalance
+        else
+            echo "irqbalance 服务已安装，但当前未运行"
+        fi
+    elif service --status-all 2>/dev/null | grep -q irqbalance; then
+        if service irqbalance status >/dev/null 2>&1; then
+            echo "检测到 irqbalance 服务正在运行，正在关闭..."
+            service irqbalance stop
+            echo "irqbalance 已关闭"
+        else
+            echo "irqbalance 服务已安装，但当前未运行"
+        fi
+    else
+        echo "当前环境未检测到 irqbalance 服务"
+    fi
+}
+
+bind_irq() {
+    echo "==== 开始对所有 NPU 卡的中断绑核 ===="
+    check_and_stop_irqbalance
+    SQ_IRQ_LIST=($(cat /proc/interrupts | grep sq_send_trigger_irq | cut -d: -f1))
+    for i in "${!all_npu_id[@]}"; do
+        SQ_CPU=${all_start_cpus[$i]}
+        CQ_CPU=$((SQ_CPU+1))
+
+        if [[ "${#all_npu_id[@]}" -eq 8 ]]; then
+            CARD=${all_npu_id[$i]}
+            CHIP_ID=0
+        else
+            CARD=$((all_npu_id[$i] / 2))
+            CHIP_ID=$((all_npu_id[$i] % 2))
+        fi
+
+        # 获取 PCI 地址
+        PCI_ADDR=$(npu-smi info -t board -i $CARD -c $CHIP_ID| grep "PCIe Bus Info" | awk '{print $NF}' | tr '[:upper:]' '[:lower:]')
+        if [ -z "$PCI_ADDR" ]; then
+            echo "未找到 NPU 卡 $CARD 的 PCI 地址"
+            continue
+        fi
+
+        # 获取该卡的 MSI 中断号列表
+        NPU_IRQ_LIST=$(ls /sys/bus/pci/devices/$PCI_ADDR/msi_irqs/ | sort -n)
+        for irq in "${SQ_IRQ_LIST[@]}"; do
+            if echo "${NPU_IRQ_LIST[@]}" | grep -qw "$irq"; then
+                SQ_IRQ=$irq
+                CQ_IRQ=$((SQ_IRQ+1))
+            fi
+        done
+        if [ -z "$SQ_IRQ" ]; then
+            echo "未找到 NPU 卡 $CARD 的 SQ 中断"
+            continue
+        fi
+
+        echo "NPU卡 ${all_npu_id[$i]} (PCI $PCI_ADDR): SQ IRQ=$SQ_IRQ → CPU$SQ_CPU, CQ IRQ=$CQ_IRQ → CPU$CQ_CPU"
+
+        # 绑 SQ
+        echo $(cpu_to_mask $SQ_CPU) > /proc/irq/$SQ_IRQ/smp_affinity
+        # 绑 CQ
+        echo $(cpu_to_mask $CQ_CPU) > /proc/irq/$CQ_IRQ/smp_affinity
+    done
+}
+
+bind_irq

xtuner/v1/data_proto/sequence_context.py

@@ -87,7 +87,9 @@ def __init__(
         # the argument can be an int, but as an attribute it can only be a tensor
         self.input_ids = input_ids
         self.cu_seq_lens_q = cu_seq_lens_q
+        self.cu_seq_lens_q_list = cu_seq_lens_q.tolist() if isinstance(cu_seq_lens_q, torch.Tensor) else cu_seq_lens_q
         self.cu_seq_lens_k = cu_seq_lens_k
+        self.cu_seq_lens_k_list = cu_seq_lens_k.tolist() if isinstance(cu_seq_lens_k, torch.Tensor) else cu_seq_lens_k
         # force max_length_q and max_length_k be cpu tensors to avoid cuda synchronization
         # max_length_q and max_length_k should be unpacked to int in attention implementation
         if isinstance(max_length_q, int):

xtuner/v1/data_proto/utils.py

@@ -121,7 +121,7 @@ def gather_for_sequence_parallel(input, dim: int, sp_group: dist.ProcessGroup):
         return input
 
     tensor_list = [torch.empty_like(input) for _ in range(world_size)]
-    assert input.device.type == "cuda"
+    assert input.device.type == "cuda" or input.device.type == "npu"
     dist.all_gather(tensor_list, input, group=sp_group)
 
     output = torch.cat(tensor_list, dim=dim).contiguous()

xtuner/v1/loss/ce_loss.py

@@ -264,6 +264,7 @@ def forward(
         hidden_states: torch.Tensor,
         head_weight: torch.Tensor,
         head_bias: torch.Tensor | None = None,
+        skip_all_reduce: bool = False,
     ) -> tuple[torch.Tensor, tuple[torch.Tensor | None, dict[str, Any]]]:
         from xtuner.v1.model.utils.misc import ModelForwardExtraLogInfo
 
@@ -282,9 +283,9 @@ def forward(
 
         extra_info["local_base_loss"] = loss.detach().clone()
 
-        # Step 2.c in the loss calculation: reduce the loss over all ranks using all_reduce with autograd support
-        if dist.is_initialized():
-            loss = all_reduce(loss, op=dist.ReduceOp.SUM, group=dist.group.WORLD)
+        if not skip_all_reduce:
+            if dist.is_initialized():
+                loss = all_reduce(loss, op=dist.ReduceOp.SUM, group=dist.group.WORLD)
 
         return loss, (logits, extra_info)
 

xtuner/v1/loss/mtp_loss.py

@@ -162,12 +162,12 @@ def forward(
         hidden_states: torch.Tensor,
         head_weight: torch.Tensor,
         head_bias: torch.Tensor | None = None,
+        skip_all_reduce: bool = False,
     ) -> tuple[torch.Tensor, tuple[torch.Tensor | None, dict[str, Any]]]:
         if self.loss_cfg.detach_mtp_lm_head_weight:
             head_weight = head_weight.detach()
             head_bias = head_bias.detach() if head_bias is not None else None
-        # Dispatch to eager_mode/chunk_mode via base class, which calls loss_fn per chunk
-        return super().forward(hidden_states, head_weight, head_bias)
+        return super().forward(hidden_states, head_weight, head_bias, skip_all_reduce=skip_all_reduce)
 
     def loss_fn(
         self,

xtuner/v1/model/compose/qwen3_vl/modeling_vision.py

@@ -503,6 +503,7 @@ def forward(self, hidden_states: torch.Tensor,
         )
         cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
         max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+        cu_seqlens_list: List[int] = cu_seqlens.tolist() if isinstance(cu_seqlens, torch.Tensor) else cu_seqlens
 
         if sequence_parallel_mesh and sequence_parallel_mesh.size() > 1:
             div_num = sequence_parallel_mesh.size() * 4
@@ -532,15 +533,15 @@ def forward(self, hidden_states: torch.Tensor,
                 ):
                     hidden_states = blk(
                         hidden_states,
-                        cu_seqlens,
+                        cu_seqlens_list,
                         max_seqlen,
                         position_embeddings,
                         sequence_parallel_mesh,
                     )
             else:
                 hidden_states = blk(
                     hidden_states,
-                    cu_seqlens,
+                    cu_seqlens_list,
                     max_seqlen,
                     position_embeddings,
                     sequence_parallel_mesh,

xtuner/v1/model/moe/moe.py

@@ -422,6 +422,7 @@ def _micro_batch_forward(
         seq_ctx_list: list[SequenceContext],
         loss_ctx_list: list[MoELossContextDict],
         return_router_logits: bool = False,
+        skip_all_reduce: bool = False,
     ) -> MoEModelOutputs:
         """Micro-batch forward pass for MoE model.
 
@@ -588,7 +589,9 @@ def _micro_batch_forward(
                 micro_batch_mtp_losses = torch.tensor(0.0, device=DEVICE)
                 for mtp_idx, (mtp_hidden, mtp_ctx) in enumerate(zip(mtp_outputs, mtp_loss_ctx_list)):
                     mtp_hidden_states, mtp_router_results, _ = mtp_hidden
-                    mtp_loss, _ = self.lm_head(mtp_hidden_states, cast(MTPLossContext, mtp_ctx))
+                    mtp_loss, _ = self.lm_head(
+                        mtp_hidden_states, cast(MTPLossContext, mtp_ctx), skip_all_reduce=skip_all_reduce
+                    )
                     micro_batch_mtp_losses += mtp_loss
 
                     if keep_router:
@@ -598,7 +601,15 @@ def _micro_batch_forward(
                 has_mtp_loss = True
 
             if has_mtp_loss:
-                output["mtp_loss"] = mtp_losses * self.config.mtp_config.loss_scaling_factor
+                if skip_all_reduce:
+                    scaled_mtp_loss = mtp_losses * self.config.mtp_config.loss_scaling_factor
+                    if dist.is_initialized() and scaled_mtp_loss.requires_grad:
+                        scaled_mtp_loss = torch.distributed.nn.functional.all_reduce(
+                            scaled_mtp_loss, op=dist.ReduceOp.SUM, group=dist.group.WORLD
+                        )
+                    output["mtp_loss"] = scaled_mtp_loss
+                else:
+                    output["mtp_loss"] = mtp_losses * self.config.mtp_config.loss_scaling_factor
 
         # Apply final norm to all micro-batches
         cat_hidden_states = torch.cat(hidden_states_list, dim=1)
@@ -650,6 +661,7 @@ def _forward(
         seq_ctx: SequenceContext,  # todo(@yehaochen): support intra layer micro-batch
         loss_ctx: MoELossContextDict | None,
         return_router_logits: bool = False,
+        skip_all_reduce: bool = False,
     ) -> MoEModelOutputs:
         input_ids = seq_ctx.input_ids
         position_ids = seq_ctx.position_ids
@@ -795,15 +807,23 @@ def _forward(
                     num_tokens_global=mtp_num_tokens_global,
                     world_size=mtp_z_world_size,
                 )
-                mtp_loss, _ = self.lm_head(mtp_hidden_states, cast(MTPLossContext, mtp_ctx))
+                mtp_loss, _ = self.lm_head(
+                    mtp_hidden_states, cast(MTPLossContext, mtp_ctx), skip_all_reduce=skip_all_reduce
+                )
                 mtp_losses += mtp_loss
 
             # Average MTP losses across depths and scale
             mtp_losses = mtp_losses / len(mtp_loss_ctx_list)
             scaled_mtp_loss = mtp_losses * self.config.mtp_config.loss_scaling_factor  # type: ignore
 
-            # Add to total loss
-            output["mtp_loss"] = scaled_mtp_loss
+            if skip_all_reduce:
+                if dist.is_initialized() and scaled_mtp_loss.requires_grad:
+                    scaled_mtp_loss = torch.distributed.nn.functional.all_reduce(
+                        scaled_mtp_loss, op=dist.ReduceOp.SUM, group=dist.group.WORLD
+                    )
+                output["mtp_loss"] = scaled_mtp_loss
+            else:
+                output["mtp_loss"] = scaled_mtp_loss
 
         split_aux_output = self.aux_loss.finalize(
             balancing_ctx=balancing_ctx,

xtuner/v1/model/moe/qwen3vl_text.py

@@ -113,9 +113,10 @@ def _forward(
         seq_ctx: SequenceContext,  # todo(@yehaochen): support intra layer micro-batch
         loss_ctx: MoELossContextDict | None,
         return_router_logits: bool = False,
+        skip_all_reduce: bool = False,
     ) -> MoEModelOutputs:
         if seq_ctx.deepstack_visual_embeds is None:
-            return super()._forward(seq_ctx, loss_ctx, return_router_logits)
+            return super()._forward(seq_ctx, loss_ctx, return_router_logits, skip_all_reduce)
 
         input_ids = seq_ctx.input_ids
         position_ids = seq_ctx.position_ids
@@ -210,7 +211,7 @@ def _forward(
 
         # Get LM loss context from dict
         lm_loss_ctx = loss_ctx["lm"] if loss_ctx is not None else None
-        loss, (logits, extra_info) = self.lm_head(hidden_states, lm_loss_ctx)  # type: ignore
+        loss, (logits, extra_info) = self.lm_head(hidden_states, lm_loss_ctx, skip_all_reduce=skip_all_reduce)  # type: ignore
         output["loss"] = loss
         output["logits"] = logits
         output["extra_info"] = extra_info

xtuner/v1/module/attention/mha.py

@@ -405,8 +405,12 @@ def forward(
             query_states,
             key_states,
             value_states,
-            cu_seqlens_q=seq_ctx.cu_seq_lens_q,
-            cu_seqlens_k=seq_ctx.cu_seq_lens_k,
+            cu_seqlens_q=seq_ctx.cu_seq_lens_q_list
+            if hasattr(seq_ctx, "cu_seq_lens_q_list")
+            else seq_ctx.cu_seq_lens_q,
+            cu_seqlens_k=seq_ctx.cu_seq_lens_k_list
+            if hasattr(seq_ctx, "cu_seq_lens_k_list")
+            else seq_ctx.cu_seq_lens_k,
             max_seqlen_q=seq_ctx.max_length_q,
             max_seqlen_k=seq_ctx.max_length_k,
             window_size=self.window_size,

xtuner/v1/module/dispatcher/torch_all2all.py

@@ -471,10 +471,11 @@ def combine_preprocess(
         async_op: bool = False,
         decoding: bool = False,
     ) -> TorchAll2AllPreCombineResult:
-        hidden_states = unpermute(
-            hidden_states,
-            post_dispatched["row_ids_map"],
-        )
+        if len(hidden_states) != 0:
+            hidden_states = unpermute(
+                hidden_states,
+                post_dispatched["row_ids_map"],
+            )
 
         if async_op:
             backward_previous_event = cast(torch.cuda.Event, torch.cuda.Event())

xtuner/v1/module/lm_head/lm_head.py

@@ -20,16 +20,16 @@
 class LMHead(nn.Linear):
     @overload  # type: ignore[override]
     def forward(
-        self, hidden_states: HiddenStates, loss_ctx: None = None
+        self, hidden_states: HiddenStates, loss_ctx: None = None, *, skip_all_reduce: bool = False
     ) -> tuple[None, tuple[Logits | None, dict[str, Any]]]: ...
 
     @overload  # type: ignore[override]
     def forward(
-        self, hidden_states: HiddenStates, loss_ctx: LMHeadLossContext
+        self, hidden_states: HiddenStates, loss_ctx: LMHeadLossContext, *, skip_all_reduce: bool = False
     ) -> tuple[Loss, tuple[Logits | None, dict[str, Any]]]: ...
 
     def forward(  # type: ignore[override]
-        self, hidden_states: torch.Tensor, loss_ctx: LMHeadLossContext | None = None
+        self, hidden_states: torch.Tensor, loss_ctx: LMHeadLossContext | None = None, *, skip_all_reduce: bool = False
     ) -> tuple[Loss | None, tuple[Logits | None, dict[str, Any]]]:
         """Forward pass of the language model head."""
         if isinstance(self.weight, DTensor):
@@ -46,16 +46,16 @@ def forward(  # type: ignore[override]
             logits = F.linear(hidden_states, w, b)
             return None, (logits.float(), {})
         else:
-            return loss_ctx.forward(hidden_states, w, b)
+            return loss_ctx.forward(hidden_states, w, b, skip_all_reduce=skip_all_reduce)
 
     @overload  # type: ignore
     def __call__(
-        self, hidden_states: HiddenStates, loss_ctx: None = None
+        self, hidden_states: HiddenStates, loss_ctx: None = None, *, skip_all_reduce: bool = False
     ) -> tuple[None, tuple[Logits | None, dict[str, Any]]]: ...
 
     @overload  # type: ignore
     def __call__(
-        self, hidden_states: HiddenStates, loss_ctx: LMHeadLossContext
+        self, hidden_states: HiddenStates, loss_ctx: LMHeadLossContext, *, skip_all_reduce: bool = False
     ) -> tuple[Loss, tuple[Logits | None, dict[str, Any]]]: ...
 
     __call__ = nn.Module.__call__

xtuner/v1/ops/flash_attn/npu.py

@@ -31,8 +31,12 @@ def npu_flash_varlen_attn(
             scale=q.shape[-1] ** -0.5,
             keep_prob=1 - dropout_p,
             input_layout="TND",
-            actual_seq_qlen=tuple(cu_seqlens_q[1:].tolist()),
-            actual_seq_kvlen=tuple(cu_seqlens_k[1:].tolist()),
+            actual_seq_qlen=tuple(cu_seqlens_q[1:].tolist())
+            if isinstance(cu_seqlens_q, torch.Tensor)
+            else cu_seqlens_q[1:],
+            actual_seq_kvlen=tuple(cu_seqlens_k[1:].tolist())
+            if isinstance(cu_seqlens_k, torch.Tensor)
+            else cu_seqlens_k[1:],
         )[0]
     else:
         fa_out = torch_npu.npu_fusion_attention(
@@ -45,8 +49,12 @@ def npu_flash_varlen_attn(
             scale=q.shape[-1] ** -0.5,
             keep_prob=1 - dropout_p,
             input_layout="TND",
-            actual_seq_qlen=tuple(cu_seqlens_q[1:].tolist()),
-            actual_seq_kvlen=tuple(cu_seqlens_k[1:].tolist()),
+            actual_seq_qlen=tuple(cu_seqlens_q[1:].tolist())
+            if isinstance(cu_seqlens_q, torch.Tensor)
+            else cu_seqlens_q[1:],
+            actual_seq_kvlen=tuple(cu_seqlens_k[1:].tolist())
+            if isinstance(cu_seqlens_k, torch.Tensor)
+            else cu_seqlens_k[1:],
             sparse_mode=3,
         )[0]
     return fa_out

xtuner/v1/ops/rms_norm/__init__.py

@@ -34,6 +34,13 @@ def _norm(x):
     return output.type_as(x)
 
 
+def zero_centered_rms_norm_npu(x: torch.Tensor, weight: torch.Tensor, epsilon: float) -> torch.Tensor:
+    import torch_npu
+
+    output = torch_npu.npu_rms_norm(x, 1.0 + weight.float(), epsilon)[0]
+    return output.type_as(x)
+
+
 def get_rms_norm_fn() -> RMSNormProtocol:
     from xtuner.v1.utils import get_device
 
@@ -63,11 +70,14 @@ def get_zero_centered_rms_norm_fn() -> RMSNormProtocol:
         else:
             return native_zero_centered_rms_norm
     elif device == "npu":
+        if os.getenv("XTUNER_USE_NATIVE_RMSNORM", "1") == "0":
+            raise NotImplementedError("Zero-centered RMSNorm is not implemented in triton")
+        else:
+            return zero_centered_rms_norm_npu
+        # def _not_implemented(*args, **kwargs):
+        #     raise NotImplementedError("Zero-centered RMSNorm is not implemented on NPU")
 
-        def _not_implemented(*args, **kwargs):
-            raise NotImplementedError("Zero-centered RMSNorm is not implemented on NPU")
-
-        return _not_implemented
+        # return _not_implemented
     else:
         raise NotImplementedError(f"RMSNorm is not implemented on {device}")