From 274bbdce978e31c990e7b317da033c15080e8cc7 Mon Sep 17 00:00:00 2001 From: wentiange Date: Mon, 13 Jul 2026 04:46:11 +0000 Subject: [PATCH 1/3] =?UTF-8?q?=E5=87=8F=E5=B0=91=E9=80=9A=E4=BF=A1?= =?UTF-8?q?=E6=AC=A1=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- xtuner/v1/data_proto/sequence_context.py | 2 ++ xtuner/v1/loss/ce_loss.py | 7 +++-- xtuner/v1/loss/mtp_loss.py | 4 +-- .../model/compose/qwen3_vl/modeling_vision.py | 5 ++-- xtuner/v1/model/moe/moe.py | 30 +++++++++++++++---- xtuner/v1/model/moe/qwen3vl_text.py | 5 ++-- xtuner/v1/module/attention/mha.py | 8 +++-- xtuner/v1/module/lm_head/lm_head.py | 12 ++++---- xtuner/v1/ops/flash_attn/npu.py | 16 +++++++--- 9 files changed, 63 insertions(+), 26 deletions(-) diff --git a/xtuner/v1/data_proto/sequence_context.py b/xtuner/v1/data_proto/sequence_context.py index 84ac054ba9..4d0d1fec32 100644 --- a/xtuner/v1/data_proto/sequence_context.py +++ b/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): diff --git a/xtuner/v1/loss/ce_loss.py b/xtuner/v1/loss/ce_loss.py index eba945fae7..703aef0e63 100644 --- a/xtuner/v1/loss/ce_loss.py +++ b/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) diff --git a/xtuner/v1/loss/mtp_loss.py b/xtuner/v1/loss/mtp_loss.py index a5aebc3b19..88aa4d1d66 100644 --- a/xtuner/v1/loss/mtp_loss.py +++ b/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, diff --git a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py index d9b599b78e..01a946e57b 100644 --- a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py +++ b/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,7 +533,7 @@ def forward(self, hidden_states: torch.Tensor, ): hidden_states = blk( hidden_states, - cu_seqlens, + cu_seqlens_list, max_seqlen, position_embeddings, sequence_parallel_mesh, @@ -540,7 +541,7 @@ def forward(self, hidden_states: torch.Tensor, else: hidden_states = blk( hidden_states, - cu_seqlens, + cu_seqlens_list, max_seqlen, position_embeddings, sequence_parallel_mesh, diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 061e2f6e29..142dc26f3f 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/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, diff --git a/xtuner/v1/model/moe/qwen3vl_text.py b/xtuner/v1/model/moe/qwen3vl_text.py index fafb70a80e..94a10ab852 100644 --- a/xtuner/v1/model/moe/qwen3vl_text.py +++ b/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 diff --git a/xtuner/v1/module/attention/mha.py b/xtuner/v1/module/attention/mha.py index 707e978055..eb4c6bc980 100644 --- a/xtuner/v1/module/attention/mha.py +++ b/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, diff --git a/xtuner/v1/module/lm_head/lm_head.py b/xtuner/v1/module/lm_head/lm_head.py index 67e912d47c..c5c49904bd 100644 --- a/xtuner/v1/module/lm_head/lm_head.py +++ b/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__ diff --git a/xtuner/v1/ops/flash_attn/npu.py b/xtuner/v1/ops/flash_attn/npu.py index 31ecb56def..b121d5e260 100644 --- a/xtuner/v1/ops/flash_attn/npu.py +++ b/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 From 6be241668d58048b060c01b3b65a7da5ca5c0c9c Mon Sep 17 00:00:00 2001 From: wentiange Date: Mon, 13 Jul 2026 04:50:09 +0000 Subject: [PATCH 2/3] =?UTF-8?q?[bugfix]=20NPU=E5=8A=9F=E8=83=BD=E6=80=A7?= =?UTF-8?q?=E9=80=82=E9=85=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- xtuner/v1/data_proto/utils.py | 2 +- xtuner/v1/module/dispatcher/torch_all2all.py | 9 +++++---- xtuner/v1/ops/rms_norm/__init__.py | 18 ++++++++++++++---- xtuner/v1/ops/rotary_emb.py | 8 ++++---- xtuner/v1/optim/muon.py | 16 +++++++++++++--- 5 files changed, 37 insertions(+), 16 deletions(-) diff --git a/xtuner/v1/data_proto/utils.py b/xtuner/v1/data_proto/utils.py index 9ceafdd06d..04a46f4cef 100644 --- a/xtuner/v1/data_proto/utils.py +++ b/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() diff --git a/xtuner/v1/module/dispatcher/torch_all2all.py b/xtuner/v1/module/dispatcher/torch_all2all.py index ba1d021e6a..7a8faa48f7 100644 --- a/xtuner/v1/module/dispatcher/torch_all2all.py +++ b/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()) diff --git a/xtuner/v1/ops/rms_norm/__init__.py b/xtuner/v1/ops/rms_norm/__init__.py index 26534e744b..93d44dbf0d 100644 --- a/xtuner/v1/ops/rms_norm/__init__.py +++ b/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}") diff --git a/xtuner/v1/ops/rotary_emb.py b/xtuner/v1/ops/rotary_emb.py index 60ea16c0fe..c4ddd5d74d 100644 --- a/xtuner/v1/ops/rotary_emb.py +++ b/xtuner/v1/ops/rotary_emb.py @@ -174,11 +174,11 @@ def get_apply_rotary_emb( device = get_device() if device == "npu": - assert not enable_partial_rotary, "Partial rotary is not supported on NPU yet." - if fope_sep_head: - return apply_rotary_pos_emb_sep_npu + assert fope_sep_head is None or not fope_sep_head, "FoPE with sep head is not supported on NPU yet." + if enable_partial_rotary: + return apply_rotary_pos_emb_cuda_for_partial_rotary else: - return apply_rotary_pos_emb_npu + return apply_rotary_pos_emb_cuda else: if fope_sep_head: logger.debug("Using FoPE with fope_sep_head") diff --git a/xtuner/v1/optim/muon.py b/xtuner/v1/optim/muon.py index fc3592024e..02f12a64b2 100644 --- a/xtuner/v1/optim/muon.py +++ b/xtuner/v1/optim/muon.py @@ -1344,7 +1344,18 @@ def adamw_update_foreach_async( epsilon: float, ) -> Generator[None, None, None]: """Async wrapper around foreach AdamW update.""" - adamw_update_foreach(X, G, M, V, lr, beta1, beta2, weight_decay, step, epsilon) + adamw_update_foreach( + X, + G, + M, + V, + lr.to(X[0].device), + beta1.to(X[0].device), + beta2.to(X[0].device), + weight_decay.to(X[0].device), + step, + epsilon, + ) yield @@ -1449,7 +1460,6 @@ def zeropower_via_newtonschulz5(G: Tensor, epsilon: float = 1e-7, num_experts: i ] X = G.to(dtype=torch.bfloat16) - original_shape = X.shape # Unified handling: reshape to (num_experts, M, N) for both cases # For regular case (num_experts=1), this adds a batch dimension of size 1 @@ -1481,6 +1491,6 @@ def zeropower_via_newtonschulz5(G: Tensor, epsilon: float = 1e-7, num_experts: i X = X.mT # Reshape back to original shape: (num_experts, M, N) -> (num_experts * M, N) - X = X.view(original_shape) + X = X.reshape(num_experts * X.size(-2), X.size(-1)) return X From bcf23ae6c1274a5c5ff3469654ae41b9cce35b20 Mon Sep 17 00:00:00 2001 From: wentiange Date: Mon, 13 Jul 2026 04:55:05 +0000 Subject: [PATCH 3/3] =?UTF-8?q?=E7=BB=91=E6=A0=B8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .dev_scripts/npu_bind_irq.sh | 89 ++++ xtuner/v1/train/trainer.py | 4 + xtuner/v1/utils/npu_cpu_binder.py | 652 ++++++++++++++++++++++++++++++ 3 files changed, 745 insertions(+) create mode 100644 .dev_scripts/npu_bind_irq.sh create mode 100644 xtuner/v1/utils/npu_cpu_binder.py diff --git a/.dev_scripts/npu_bind_irq.sh b/.dev_scripts/npu_bind_irq.sh new file mode 100644 index 0000000000..0327da445c --- /dev/null +++ b/.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 \ No newline at end of file diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py index e61f46b224..a255590847 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -2024,7 +2024,11 @@ def _resolve_load_checkpoint_cfg( def _load_checkpoint(self): load_checkpoint_cfg: LoadCheckpointConfig = self._load_checkpoint_cfg + # 从 xtuner.v1.utils 获取 npu_cpu_binder 模块 + from xtuner.v1.utils import npu_cpu_binder + # 执行npu_cpu_binder 细粒度绑核操作 + npu_cpu_binder.run(self.rank) if (resume_from := load_checkpoint_cfg.checkpoint_path) is None: log_rank0.info("No checkpoint to resume from.") return diff --git a/xtuner/v1/utils/npu_cpu_binder.py b/xtuner/v1/utils/npu_cpu_binder.py new file mode 100644 index 0000000000..95cc187ae6 --- /dev/null +++ b/xtuner/v1/utils/npu_cpu_binder.py @@ -0,0 +1,652 @@ +import logging +import os +import re +import shutil +import subprocess +from collections import defaultdict +from typing import Dict, List, Optional, Tuple + +import psutil + + +CPU_MASK_BIT = 32 +MAIN_PROCESS_RANGE = 5 +ACL_THREAD_RANGE = 1 +RELEASE_THREAD_RANGE = 1 +ALLOWED_CPUS_PATH = "/proc/self/status" + +DEFAULT_CPU_BIND_CONFIG = { + "custom_bind": [ + { + "process_name": "xtuner.v1.train.cli.sft", + "cpu_list": [ + "15-22", + "29-36", + "55-62", + "69-76", + "95-102", + "109-116", + "135-142", + "149-156", + "175-182", + "189-196", + "215-222", + "229-236", + "255-262", + "269-276", + "295-302", + "309-316", + ], + "bind_sub_process": True, + }, + { + "process_name": "acl_thread", + "cpu_list": [ + "23", + "37", + "63", + "77", + "103", + "117", + "143", + "157", + "183", + "197", + "223", + "237", + "263", + "277", + "303", + "317", + ], + "is_thread": True, + "mem_bind": True, + }, + { + "process_name": "release_thread", + "cpu_list": [ + "24", + "38", + "64", + "78", + "104", + "118", + "144", + "158", + "184", + "198", + "224", + "238", + "264", + "278", + "304", + "318", + ], + "is_thread": True, + }, + { + "process_name": "hccp_connect", + "cpu_list": [ + "25", + "39", + "65", + "79", + "105", + "119", + "145", + "159", + "185", + "199", + "225", + "239", + "265", + "279", + "305", + "319", + ], + "is_thread": True, + }, + ] +} + +logging.basicConfig(level=logging.INFO, format="[%(asctime)s] [%(levelname)s]:%(message)s") + + +def execute_command(cmd: List[str]) -> Tuple[str, int]: + with subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE) as p: + out, _ = p.communicate(timeout=1000) + return out.decode(), p.returncode + + +def cpu_to_mask(cpus: List[int]) -> str: + groups: Dict[int, int] = defaultdict(int) + for cpu in cpus: + group = cpu // CPU_MASK_BIT + bit = cpu % CPU_MASK_BIT + groups[group] |= 1 << bit + + max_group = max(groups.keys()) + mask_parts = [] + for group in reversed(range(max_group + 1)): + mask_parts.append(f"{groups.get(group, 0):08x}") + return ",".join(mask_parts) + + +def expand_cpu_list(cpu_str: str) -> List[int]: + cpus: List[int] = [] + for part in cpu_str.split(","): + if "-" in part: + start, end = map(int, part.split("-")) + cpus.extend(range(start, end + 1)) + else: + cpus.append(int(part)) + return cpus + + +def get_npu_map_info() -> Dict[str, Dict[str, str]]: + npu_map_info: Dict[str, Dict[str, str]] = {} + npu_info, _ = execute_command(["npu-smi", "info", "-m"]) + npu_map = npu_info.strip().split("\n")[1:] + for line in npu_map: + parts = line.strip().split() + if len(parts) < 3: + continue + npu_id, chip_id, chip_logic_id = parts[:3] + if chip_logic_id.isdigit(): + npu_map_info.setdefault(npu_id, {})[chip_id] = chip_logic_id + logging.debug(f"build npu_map_info: {npu_map_info}") + return npu_map_info + + +class DeviceInfo: + def __init__(self): + self.main_pid_list: List[List[int]] = [] + self.npu_map_info: Dict[str, Dict[str, str]] = get_npu_map_info() + self.allowed_cpus: List[int] = self.parse_allowed_cpus() + self.running_npu_list: List[int] = self.get_running_npus() + self.npu_affinity: Dict[int, List[int]] = self.parse_topo_affinity() + + @staticmethod + def parse_allowed_cpus() -> List[int]: + if not os.path.exists(ALLOWED_CPUS_PATH): + return [] + with open(ALLOWED_CPUS_PATH) as f: + for line in f: + if line.startswith("Cpus_allowed_list"): + allowed_cpu_list = expand_cpu_list(line.split()[1]) + logging.debug(f"CPUs_allowed_list: {allowed_cpu_list}") + return allowed_cpu_list + return [] + + @staticmethod + def parse_topo_affinity() -> Dict[int, List[int]]: + chip_logic_id = 0 + affinity: Dict[int, List[int]] = {} + affinity_message, _ = execute_command(["npu-smi", "info", "-t", "topo"]) + for line in affinity_message.splitlines(): + if line.startswith("NPU"): + parts = line.split() + last_part = parts[-1] + if last_part != "Affinity": + affinity[chip_logic_id] = expand_cpu_list(last_part) + chip_logic_id += 1 + logging.debug(f"build affinity map: {affinity}") + return affinity + + def get_running_npus(self) -> List[int]: + npu_message, _ = execute_command(["npu-smi", "info"]) + in_proc_section = False + running_npu_set: set[int] = set() + chip_pid_map: Dict[int, List[Tuple[int, int]]] = {} + for line in npu_message.splitlines(): + line = line.strip() + if line.startswith("| NPU") and "Process id" in line: + in_proc_section = True + continue + if not in_proc_section or not line.startswith("| "): + continue + parts = [p.strip() for p in line.strip("|").split("|")] + if len(parts) < 4 or not parts[1].isdigit(): + continue + pid = int(parts[1]) + try: + mem = int(parts[3]) + except ValueError: + mem = 0 + npu_id, chip_id = parts[0].split()[:2] + chip_logic_id_str = self.npu_map_info.get(npu_id, {}).get(chip_id) + if chip_logic_id_str and chip_logic_id_str.isdigit(): + chip_logic_id = int(chip_logic_id_str) + chip_pid_map.setdefault(chip_logic_id, []).append((pid, mem)) + running_npu_set.add(chip_logic_id) + + self.main_pid_list = [] + for npu in sorted(running_npu_set): + pid_mem_list: List[Tuple[int, int]] = chip_pid_map.get(npu, []) + if pid_mem_list: + max_pid = max(pid_mem_list, key=lambda x: x[1])[0] + self.main_pid_list.append([max_pid]) + logging.debug(f"identifying the running NPU card: {running_npu_set}") + return sorted(running_npu_set) + + +class CpuAlloc: + def __init__(self, device_info: Optional[DeviceInfo] = None): + self.device_info: DeviceInfo = device_info or DeviceInfo() + self.cpu_node: Dict[int, int] = {} + self.numa_to_cpu_map: Dict[int, List[int]] = defaultdict(list) + self.npu_cpu_pool: Dict[int, List[int]] = {} + self.npu_cpu_pool_all: Dict[int, List[int]] = {} + self.assign_main: Dict[int, List[int]] = {} + self.assign_acl: Dict[int, List[int]] = {} + self.assign_rel: Dict[int, List[int]] = {} + + @staticmethod + def average_distribute(groups: Dict[str, List[int]], pool: Dict[int, List[int]]) -> Dict[int, List[int]]: + result: Dict[int, List[int]] = {} + for key, npu_list in groups.items(): + cpu_list = sorted(pool[npu_list[0]]) + cpu_num_per_npu = len(cpu_list) // len(npu_list) + for i, npu in enumerate(npu_list): + start_index = i * cpu_num_per_npu + end_index = (i + 1) * cpu_num_per_npu if i < len(npu_list) - 1 else len(cpu_list) + result[npu] = cpu_list[start_index:end_index] + return result + + @staticmethod + def get_acl_main_threads() -> List[int]: + thread_message, _ = execute_command(["ps", "-Te"]) + pids: List[int] = [] + acl_threads_set = set() + for line in thread_message.splitlines(): + if "acl_thread" in line: + pid = line.split()[0] + if pid not in acl_threads_set: + acl_threads_set.add(pid) + pids.append(int(pid)) + return pids + + def dev_alloc(self) -> tuple[List[int], List[str]]: + dev_pid_list: List[int] = [] + dev_cpu_list: List[str] = [] + out, _ = execute_command(["ps", "aux"]) + for line in out.splitlines(): + m = re.search(r"dev(\d+)_sq_task", line) + if not m: + continue + dev_id = int(m.group(1)) + pid = int(line.split()[1]) + cpus = self.npu_cpu_pool_all.get(dev_id, []) + if cpus: + core = cpus[2] if len(cpus) >= 3 else cpus[0] + dev_pid_list.append(pid) + dev_cpu_list.append(str(core)) + return dev_pid_list, dev_cpu_list + + def irq_alloc(self) -> tuple[List[int], List[str]]: + sq_irqs = [] + irq_id_list: List[int] = [] + irq_cpu_list: List[str] = [] + with open("/proc/interrupts") as f: + for line in f: + if "sq_send_trigger_irq" in line: + irq = line.split(":")[0].strip() + sq_irqs.append(irq) + + for npu in sorted(self.npu_cpu_pool_all.keys()): + cpus = self.npu_cpu_pool_all[npu] + if len(cpus) < 2: + continue + + info, _ = execute_command(["npu-smi", "info", "-t", "board", "-i", str(npu)]) + pci_addr = "" + for line in info.splitlines(): + if "PCIe Bus Info" in line: + pci_addr = line.split()[-1].lower() + break + if not pci_addr: + raise RuntimeError(f"Can't find PCI address of NPU{npu} .") + + msi_irq_dir = f"/sys/bus/pci/devices/{pci_addr}/msi_irqs/" + if not os.path.exists(msi_irq_dir): + raise RuntimeError(f"Can't find MSI interrupt directory of NPU{npu} .") + + npu_irq_list = sorted(os.listdir(f"/sys/bus/pci/devices/{pci_addr}/msi_irqs/"), key=lambda x: int(x)) + for irq in sq_irqs: + if irq in npu_irq_list: + irq_id_list.extend([int(irq), int(irq) + 1]) + irq_cpu_list.extend([str(cpus[0]), str(cpus[1])]) + break + return irq_id_list, irq_cpu_list + + def build_cpu_node_map(self) -> None: + cpu_numa_map, _ = execute_command(["lscpu", "-e=CPU,NODE"]) + for line in cpu_numa_map.splitlines(): + line = line.strip() + if not line or not line[0].isdigit(): + continue + cpu_str, node_str = line.split() + cpu = int(cpu_str) + node = int(node_str) + self.cpu_node[cpu] = node + self.numa_to_cpu_map[node].append(cpu) + if not self.numa_to_cpu_map: + raise RuntimeError("lscpu 输出错误，未检测到 NUMA 节点") + + def extend_numa(self, cpu_list: List[int]) -> List[int]: + if not cpu_list: + return [] + nodes = {self.cpu_node[c] for c in cpu_list} + if len(nodes) != 1: + return cpu_list + node = list(nodes)[0] + next_node = (node + 1) % len(self.numa_to_cpu_map) + extended = cpu_list[:] + for cpu in self.numa_to_cpu_map[next_node]: + if cpu in self.device_info.allowed_cpus: + extended.append(cpu) + return sorted(set(extended)) + + def handle_no_affinity(self) -> None: + num_running_npu = len(self.device_info.running_npu_list) + num_numa_node = len(self.numa_to_cpu_map) + if num_numa_node == 0 or num_running_npu == 0: + return + npu_num_per_node = (num_running_npu // num_numa_node) + (1 if num_running_npu % num_numa_node else 0) + index = 0 + for node in sorted(self.numa_to_cpu_map): + cpus = [c for c in self.numa_to_cpu_map[node] if c in self.device_info.allowed_cpus] + if not cpus: + continue + npu_num_this_node = min(npu_num_per_node, num_running_npu - index) + total_cpu_num = len(cpus) + base_cpu_num = total_cpu_num // npu_num_this_node + extra_cpu_num = total_cpu_num % npu_num_this_node + start_index = 0 + for i in range(npu_num_this_node): + take_cpu_num = base_cpu_num + (1 if i < extra_cpu_num else 0) + end_index = start_index + take_cpu_num + select_cpus_list = cpus[start_index:end_index] + if index < num_running_npu: + npu = self.device_info.running_npu_list[index] + self.npu_cpu_pool[npu] = select_cpus_list + index += 1 + start_index = end_index + + def build_cpu_pools_all(self) -> None: + raw_pool: Dict[int, List[int]] = {} + + if self.device_info.npu_affinity: + for npu, cpus in self.device_info.npu_affinity.items(): + filtered = [c for c in cpus if c in self.device_info.allowed_cpus] + raw_pool[npu] = filtered + else: + self.handle_no_affinity() + raw_pool = self.npu_cpu_pool.copy() + + groups: Dict[str, List[int]] = defaultdict(list) + for npu, cpus in raw_pool.items(): + groups[str(cpus)].append(npu) + + final_pool: Dict[int, List[int]] = {} + for key, npu_list in groups.items(): + if len(npu_list) == 1: + final_pool[npu_list[0]] = raw_pool[npu_list[0]] + else: + final_pool.update(self.average_distribute({key: npu_list}, raw_pool)) + logging.debug(f"npu_cpu_pool_all: {final_pool}") + self.npu_cpu_pool_all = final_pool + + def build_cpu_pools_running(self) -> None: + self.build_cpu_node_map() + raw_pool: Dict[int, List[int]] = {} + + if self.device_info.npu_affinity: + for npu in self.device_info.running_npu_list: + cpus = self.device_info.npu_affinity.get(npu, []) + filtered = [c for c in cpus if c in self.device_info.allowed_cpus] + raw_pool[npu] = filtered + else: + self.handle_no_affinity() + for npu in self.device_info.running_npu_list: + if npu in self.npu_cpu_pool: + raw_pool[npu] = self.npu_cpu_pool[npu] + + groups: Dict[str, List[int]] = defaultdict(list) + for npu, cpus in raw_pool.items(): + groups[str(cpus)].append(npu) + + final_pool: Dict[int, List[int]] = {} + for key, npu_list in groups.items(): + if len(npu_list) == 1: + final_pool[npu_list[0]] = raw_pool[npu_list[0]] + else: + final_pool.update(self.average_distribute({key: npu_list}, raw_pool)) + logging.debug(f"npu_cpu_pool: {final_pool}") + self.npu_cpu_pool = final_pool + + def allocate(self, main_range: int, acl_range: int, rel_range: int) -> None: + for npu, pool in self.npu_cpu_pool.items(): + usable_pool = pool[3:] + need = main_range + acl_range + rel_range + if len(usable_pool) < need: + raise RuntimeError(f"NPU{npu} 的 CPU 数量不足，默认方案需要至少 {need} 个") + self.assign_main[npu] = usable_pool[:main_range] + self.assign_acl[npu] = usable_pool[main_range : main_range + acl_range] + self.assign_rel[npu] = usable_pool[main_range + acl_range : main_range + acl_range + rel_range] + + +class CustomBind: + def __init__( + self, + process_name: str = "", + cpu_list: Optional[List[str]] = None, + bind_sub_process: bool = False, + is_thread: bool = False, + is_irq: bool = False, + mem_bind: bool = False, + pid: Optional[List[int]] = None, + irq_id: Optional[List[int]] = None, + ): + self.process_name = process_name + self.bind_sub_process = bind_sub_process + self.is_thread = is_thread + self.is_irq = is_irq + self.mem_bind = mem_bind + self.pid = pid or [] + self.irq_id = irq_id or [] + self.cpu_list = [expand_cpu_list(seg) for seg in (cpu_list or [])] + + @staticmethod + def get_main_pid_from_docker(pid: int) -> int: + pid_file = f"/proc/{pid}/status" + if not os.path.exists(pid_file): + return 0 + out, return_code = execute_command(["grep", "Ngid", pid_file]) + if return_code != 0: + return 0 + parts = out.strip().split() + if parts[-1] != "0": + return int(parts[-1]) + else: + return 0 + + def get_real_main_pid_list( + self, pid_list: List[Tuple[int, int]], main_pid_list: List[List[int]] + ) -> List[List[int]]: + real_main_pid_list: List[List[int]] = [] + for pid, ppid in pid_list: + per_real_pid_list: List[int] = [] + for pids in main_pid_list: + if pid in pids: + per_real_pid_list.append(pid) + continue + elif ppid in pids: + per_real_pid_list.append(ppid) + continue + real_pid = self.get_main_pid_from_docker(pid) + if real_pid in pids: + per_real_pid_list.append(pid) + continue + real_ppid = self.get_main_pid_from_docker(ppid) + if real_ppid in pids: + per_real_pid_list.append(ppid) + if per_real_pid_list: + real_main_pid_list.append(per_real_pid_list) + return real_main_pid_list + + def find_threads(self, current_pid) -> List[Tuple[int, int]]: + if self.pid: + pid_list = [] + for p in self.pid: + try: + ppid = int(subprocess.check_output(["ps", "-o", "ppid=", "-p", str(p)], text=True).strip()) + except subprocess.CalledProcessError: + ppid = -1 + if p == current_pid or ppid == current_pid: + pid_list.append((p, ppid)) + return pid_list + + select_idx = 1 if self.is_thread else 0 + out, _ = execute_command(["ps", "-Te"]) if self.is_thread else execute_command(["ps", "-eo", "pid,ppid,cmd"]) + pid_list = [] + for line in out.splitlines(): + if self.process_name in line: + parts = line.split() + if len(parts) >= 2 and parts[0].isdigit() and parts[1].isdigit(): + pid = int(parts[select_idx]) + ppid = int(parts[1 - select_idx]) + if pid == current_pid or ppid == current_pid: + pid_list.append((pid, ppid)) + if not pid_list: + raise RuntimeError(f"未找到进程名称包含 {self.process_name} 的进程") + return pid_list + + def irq_bind(self) -> None: + if not shutil.which("systemctl"): + logging.warning( + "The systemctl command cannot be used in the current environment.If the irqbalance " + "service is enabled, manually disable the irqbalance service.Otherwise, the " + "interrupt-core binding cannot take effect." + ) + else: + out, return_code = execute_command(["systemctl", "list-unit-files"]) + if return_code == 0 and "irqbalance.service" in out: + _, return_code = execute_command(["systemctl", "is-active", "--quiet", "irqbalance"]) + if return_code == 0: + logging.info("The irqbalance service is running and has been stopped.") + _, return_code = execute_command(["systemctl", "stop", "irqbalance"]) + if return_code != 0: + logging.warning( + "The irqbalance service cannot be stopped.You need to manually stop it." + "Otherwise, the interrupt-core binding cannot take effect." + ) + + for irq_id, target_cpu_list in zip(self.irq_id, self.cpu_list): + affinity_file_path = f"/proc/irq/{irq_id}/smp_affinity" + with open(affinity_file_path, "w") as f: + f.write(cpu_to_mask(target_cpu_list)) + logging.info(f"Bind the interrupt of IRQ-{irq_id} to CPU{target_cpu_list}") + + def execute_bind(self, pid: int, cpu_list_str: str, process_type: str, cpu_node: Dict[int, int]) -> None: + cmd = ["taskset", "-acp" if self.bind_sub_process else "-cp", cpu_list_str, str(pid)] + logging.info(f"Bind the {self.process_name or 'target'} ({process_type}={pid}) to CPU{cpu_list_str}") + _, return_code = execute_command(cmd) + if return_code != 0: + raise RuntimeError(f"Failed to execute the command: {' '.join(cmd)}") + + def bind(self, cpu_allocer: CpuAlloc, current_pid: int) -> None: + process_type = "pid" if not self.is_thread else "tid" + pid_list = self.find_threads(current_pid) + if not pid_list: + return + real_main_pid_list = self.get_real_main_pid_list(pid_list, cpu_allocer.device_info.main_pid_list) + cpu_index = -1 + for pid, ppid in pid_list: + cpu_list_str = "" + if len(self.cpu_list) == 1: + cpu_list_str = ",".join(map(str, self.cpu_list[0])) + self.execute_bind(pid, cpu_list_str, process_type, cpu_allocer.cpu_node) + continue + if len(pid_list) == len(self.cpu_list): + cpu_index += 1 + cpu_list_str = ",".join(map(str, self.cpu_list[cpu_index])) + for pids in real_main_pid_list: + if pid in pids or ppid in pids: + cpu_list_str = ",".join(map(str, self.cpu_list[real_main_pid_list.index(pids)])) + if not cpu_list_str: + logging.error( + f"Failed to bind process {pid_list} to CPU {self.cpu_list}. Please ensure that " + f"the number of processes to be bound is the same as the number of cpu_list you have " + f"entered, or ensure that the Ngid field in the /proc/{pid}/status file is not 0. " + f"It is recommended that the script be executed on the host machine." + ) + continue + self.execute_bind(pid, cpu_list_str, process_type, cpu_allocer.cpu_node) + logging.debug(f"pid_list: {pid_list}") + logging.debug(f"real_main_pid_list: {real_main_pid_list}") + + +def export_bind_config(cpu_alloc: CpuAlloc, config: dict) -> Dict: + main_cpu_list: List[str] = [] + acl_cpu_list: List[str] = [] + rel_cpu_list: List[str] = [] + config = {"custom_bind": []} + cpu_alloc.build_cpu_pools_all() + cpu_alloc.allocate(MAIN_PROCESS_RANGE, ACL_THREAD_RANGE, RELEASE_THREAD_RANGE) + main_pid_list = cpu_alloc.get_acl_main_threads() + dev_pid_list, dev_cpu_list = cpu_alloc.dev_alloc() + irq_id_list, irq_cpu_list = cpu_alloc.irq_alloc() + for npu in sorted(cpu_alloc.device_info.running_npu_list): + main_cpu_list.append(",".join(map(str, cpu_alloc.assign_main[npu]))) + acl_cpu_list.append(",".join(map(str, cpu_alloc.assign_acl[npu]))) + rel_cpu_list.append(",".join(map(str, cpu_alloc.assign_rel[npu]))) + + config["custom_bind"].append({"pid": main_pid_list, "cpu_list": main_cpu_list, "bind_sub_process": True}) + config["custom_bind"].append({"process_name": "acl_thread", "cpu_list": acl_cpu_list, "is_thread": True}) + config["custom_bind"].append({"process_name": "release_thread", "cpu_list": rel_cpu_list, "is_thread": True}) + config["custom_bind"].append({"pid": dev_pid_list, "cpu_list": dev_cpu_list}) + config["custom_bind"].append({"irq_id": irq_id_list, "cpu_list": irq_cpu_list, "is_irq": True}) + return config + + +def load_custom_bind(data: Dict) -> List[CustomBind]: + binders = [] + for item in data.get("custom_bind", []): + binders.append( + CustomBind( + process_name=item.get("process_name", ""), + cpu_list=item["cpu_list"], + bind_sub_process=item.get("bind_sub_process", False), + is_thread=item.get("is_thread", False), + is_irq=item.get("is_irq", False), + mem_bind=item.get("mem_bind", False), + pid=item.get("pid", []), + irq_id=item.get("irq_id", []), + ) + ) + return binders + + +def run(rank_id: int, config: dict | None = None) -> None: + loop_count = 0 + input_data = config if config is not None else DEFAULT_CPU_BIND_CONFIG + cpu_allocer = CpuAlloc(DeviceInfo()) + current_pid = psutil.Process().pid + binder_list = load_custom_bind(input_data) + for bind in binder_list: + loop_count += 1 + logging.info(f"Start binding core round {loop_count}: {bind.__dict__}") + try: + if bind.is_irq: + bind.irq_bind() + else: + if not bind.pid and not bind.process_name: + logging.error("No input bound object. One of 'pid, process_name, irq_id' are required.") + continue + if len(bind.cpu_list) > 1: + bind.cpu_list = [bind.cpu_list[rank_id % 16]] + bind.bind(cpu_allocer, current_pid) + except Exception as e: + logging.error(e) + logging.info(f"===== Round {loop_count} of core binding has ended. =====")