torch_npu是支持多机多卡broadcast数据的,这块需要根据具体的报错来确认是什么样的情况。
另外的话,可以尝试下使用非ranktable(协商建链,不使用RANK_TABLE_FILE)的方式横向比对下。


@wangchao430 非 ranktable 的例子有吗?谢谢!


@wangchao430 非 ranktable 的例子有吗?谢谢!
例子一样的,只需要不使用RANK_TABLE_FILE这个环境变量就可以,这个时候就会走协商建链。 在脚本层面是不感知具体的建链方式的, torch_npu通过环境变量走不同的建链方式。


torch_npu是支持多机多卡broadcast数据的,这块需要根据具体的报错来确认是什么样的情况。
另外的话,可以尝试下使用非ranktable(协商建链,不使用RANK_TABLE_FILE)的方式横向比对下。
更多信息:
容器启动命令:
model_path="/data1/models/Qwen3-0.6B"
PORT="9522"
# Distributed training config
RANK_TABLE_FILE="/data1/MindIE_Auto_Deploy_Tool/rank_table_file.json"
master_addr="10.17.44.64"
master_port="12345"
node_rank=1
# ATB/HCCL config
ATB_LLM_HCCL_ENABLE=1
ATB_LLM_COMM_BACKEND="hccl"
HCCL_CONNECT_TIMEOUT=7200
HCCL_EXEC_TIMEOUT=0
ATB_LLM_BENCHMARK_ENABLE=1
ATB_LLM_ENABLE_AUTO_TRANSPOSE=0
ASCEND_GLOBAL_LOG_LEVEL=3
ASCEND_SLOG_PRINT_TO_STDOUT=1
HCCL_DETERMINISTIC=false
LCCL_DETERMINISTIC=0
HCCL_BUFFSIZE=120
ATB_WORKSPACE_MEM_ALLOC_GLOBAL=1
docker run -it \
--name dgi-dev-temp \
--network=host \
--privileged=true \
--device=/dev/davinci_manager \
--device=/dev/hisi_hdc \
--device=/dev/devmm_svm \
-v /usr/local/lib/node_modules/:/usr/local/lib/node_modules/:ro \
-v /usr/local/bin/node:/usr/local/bin/node:ro \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64/common \
-v /usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/driver \
-v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v /etc/vnpu.cfg:/etc/vnpu.cfg \
-v /usr/local/sbin:/usr/local/sbin \
-v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \
-v /data1/:/data1/ \
-v /data2/:/data2/ \
-w /root \
-u root \
--shm-size=64GB \
--ulimit msgqueue=-1 \
-e ASCEND_RT_VISIBLE_DEVICES="0,1,2,3,4,5,6,7" \
-e model_path="$model_path" \
-e host="0.0.0.0" \
-e port="$PORT" \
-e tensor_parallel_size="8" \
-e max_batch_size="-1" \
-e max_model_len="32768" \
-e max_num_seqs="128" \
-e DO_SAMPLE="false" \
-e backend="auto" \
-e npu_memory_utilization="0.8" \
-e image_save_dir="/data1/tmp" \
-e LOG_LEVEL="INFO" \
-e RANK_TABLE_FILE="$RANK_TABLE_FILE" \
-e master_addr="$master_addr" \
-e master_port="$master_port" \
-e nnodes="$nnodes" \
-e node_rank="$node_rank" \
-e ATB_LLM_HCCL_ENABLE=$ATB_LLM_HCCL_ENABLE \
-e ATB_LLM_COMM_BACKEND="$ATB_LLM_COMM_BACKEND" \
-e HCCL_CONNECT_TIMEOUT=$HCCL_CONNECT_TIMEOUT \
-e HCCL_EXEC_TIMEOUT=$HCCL_EXEC_TIMEOUT \
-e ATB_LLM_BENCHMARK_ENABLE=$ATB_LLM_BENCHMARK_ENABLE \
-e ATB_LLM_ENABLE_AUTO_TRANSPOSE=$ATB_LLM_ENABLE_AUTO_TRANSPOSE \
-e ASCEND_GLOBAL_LOG_LEVEL=$ASCEND_GLOBAL_LOG_LEVEL \
-e ASCEND_SLOG_PRINT_TO_STDOUT=$ASCEND_SLOG_PRINT_TO_STDOUT \
-e HCCL_DETERMINISTIC=$HCCL_DETERMINISTIC \
-e LCCL_DETERMINISTIC=$LCCL_DETERMINISTIC \
-e HCCL_BUFFSIZE=$HCCL_BUFFSIZE \
-e ATB_WORKSPACE_MEM_ALLOC_GLOBAL=$ATB_WORKSPACE_MEM_ALLOC_GLOBAL \
--entrypoint bash \
xxxx(可以看作cann镜像)
source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh
source venvs/dgi/bin/activate
主机:
torchrun \
--nnodes=2 \
--nproc_per_node=8 \
--node_rank=0 \
--master_addr=10.17.44.64 \
--master_port=29500 \
broadcast_object_hccl_fixed.py > /data1/master_log.txt
从机:
torchrun \
--nnodes=2 \
--nproc_per_node=8 \
--node_rank=1 \
--master_addr=10.17.44.64 \
--master_port=29500 \
broadcast_object_hccl_fixed.py > /data1/slave_log.txt
容器内运行torchrun之后的报错信息:
fe339d9b57414dc3a23599a4c54ec7b9.txt197c015d293b45e4a223dcef1b508362.txt


torch_npu是支持多机多卡broadcast数据的,这块需要根据具体的报错来确认是什么样的情况。
另外的话,可以尝试下使用非ranktable(协商建链,不使用RANK_TABLE_FILE)的方式横向比对下。
更多信息:
容器启动命令:model_path="/data1/models/Qwen3-0.6B" PORT="9522" # Distributed training config RANK_TABLE_FILE="/data1/MindIE_Auto_Deploy_Tool/rank_table_file.json" master_addr="10.17.44.64" master_port="12345" node_rank=1 # ATB/HCCL config ATB_LLM_HCCL_ENABLE=1 ATB_LLM_COMM_BACKEND="hccl" HCCL_CONNECT_TIMEOUT=7200 HCCL_EXEC_TIMEOUT=0 ATB_LLM_BENCHMARK_ENABLE=1 ATB_LLM_ENABLE_AUTO_TRANSPOSE=0 ASCEND_GLOBAL_LOG_LEVEL=3 ASCEND_SLOG_PRINT_TO_STDOUT=1 HCCL_DETERMINISTIC=false LCCL_DETERMINISTIC=0 HCCL_BUFFSIZE=120 ATB_WORKSPACE_MEM_ALLOC_GLOBAL=1 docker run -it \ --name dgi-dev-temp \ --network=host \ --privileged=true \ --device=/dev/davinci_manager \ --device=/dev/hisi_hdc \ --device=/dev/devmm_svm \ -v /usr/local/lib/node_modules/:/usr/local/lib/node_modules/:ro \ -v /usr/local/bin/node:/usr/local/bin/node:ro \ -v /usr/local/dcmi:/usr/local/dcmi \ -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ -v /usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64/common \ -v /usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/driver \ -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ -v /etc/ascend_install.info:/etc/ascend_install.info \ -v /etc/vnpu.cfg:/etc/vnpu.cfg \ -v /usr/local/sbin:/usr/local/sbin \ -v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \ -v /data1/:/data1/ \ -v /data2/:/data2/ \ -w /root \ -u root \ --shm-size=64GB \ --ulimit msgqueue=-1 \ -e ASCEND_RT_VISIBLE_DEVICES="0,1,2,3,4,5,6,7" \ -e model_path="$model_path" \ -e host="0.0.0.0" \ -e port="$PORT" \ -e tensor_parallel_size="8" \ -e max_batch_size="-1" \ -e max_model_len="32768" \ -e max_num_seqs="128" \ -e DO_SAMPLE="false" \ -e backend="auto" \ -e npu_memory_utilization="0.8" \ -e image_save_dir="/data1/tmp" \ -e LOG_LEVEL="INFO" \ -e RANK_TABLE_FILE="$RANK_TABLE_FILE" \ -e master_addr="$master_addr" \ -e master_port="$master_port" \ -e nnodes="$nnodes" \ -e node_rank="$node_rank" \ -e ATB_LLM_HCCL_ENABLE=$ATB_LLM_HCCL_ENABLE \ -e ATB_LLM_COMM_BACKEND="$ATB_LLM_COMM_BACKEND" \ -e HCCL_CONNECT_TIMEOUT=$HCCL_CONNECT_TIMEOUT \ -e HCCL_EXEC_TIMEOUT=$HCCL_EXEC_TIMEOUT \ -e ATB_LLM_BENCHMARK_ENABLE=$ATB_LLM_BENCHMARK_ENABLE \ -e ATB_LLM_ENABLE_AUTO_TRANSPOSE=$ATB_LLM_ENABLE_AUTO_TRANSPOSE \ -e ASCEND_GLOBAL_LOG_LEVEL=$ASCEND_GLOBAL_LOG_LEVEL \ -e ASCEND_SLOG_PRINT_TO_STDOUT=$ASCEND_SLOG_PRINT_TO_STDOUT \ -e HCCL_DETERMINISTIC=$HCCL_DETERMINISTIC \ -e LCCL_DETERMINISTIC=$LCCL_DETERMINISTIC \ -e HCCL_BUFFSIZE=$HCCL_BUFFSIZE \ -e ATB_WORKSPACE_MEM_ALLOC_GLOBAL=$ATB_WORKSPACE_MEM_ALLOC_GLOBAL \ --entrypoint bash \ xxxx(可以看作cann镜像)source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source venvs/dgi/bin/activate 主机: torchrun \ --nnodes=2 \ --nproc_per_node=8 \ --node_rank=0 \ --master_addr=10.17.44.64 \ --master_port=29500 \ broadcast_object_hccl_fixed.py > /data1/master_log.txt 从机: torchrun \ --nnodes=2 \ --nproc_per_node=8 \ --node_rank=1 \ --master_addr=10.17.44.64 \ --master_port=29500 \ broadcast_object_hccl_fixed.py > /data1/slave_log.txt容器内运行torchrun之后的报错信息:
fe339d9b57414dc3a23599a4c54ec7b9.txt197c015d293b45e4a223dcef1b508362.txt
@gbdjxgp 这个日志看好像没有报错? 打屏看下具体报什么样的错呢? 还有plog可以搜下ERROR发出来报错信息
另外不使用RANK_TABLE_FILE也可以试下


好的,我们这边先排查下具体的错误原因,有进展再跟您反馈!


@wangchao430 您好,我验证了下,我上次发您的程序确实是正常的,后来我尝试在mindie中的atb-models中使用dist.broadcast同步数据,结果发生报错/卡死。下面是复现过程,如果可以的话希望您能帮忙分析下!下面附上了详细的过程、时间、运行结果等!
准备文件:run_pa.py
相比源文件,添加了dist.init_process_group过程与test_infer、test_broadcast_data两个函数以及在main中调用。
# Copyright Huawei Technologies Co., Ltd. 2023-2024. All rights reserved.
import copy
import json
import math
import os
import time
import torch.distributed as dist
import torch
import torch_npu
from atb_llm.runner.model_runner import ModelRunner
from atb_llm.utils.cpu_binding import NpuHbmInfo
from atb_llm.utils.env import ENV
from atb_llm.utils.log import logger, print_log
from atb_llm.utils import file_utils
from atb_llm.models.deepseekv2.eplb.eplb_planner.eplb_worker import EplbWorker
from atb_llm.utils.moe_utils import EPLBType
from atb_llm.utils.argument_utils import ArgumentParser, StringArgumentValidator, \
NumberArgumentValidator, MAX_JSON_LENGTH
from examples.server.cache import CacheConfig, ModelConfig, CacheManager, HiddenStateCache
from examples.server.generate import decode_token, generate_req
from examples.server.request import request_from_token
OPERATOR_BOUND = 576
MAX_KW_ARGS_LENGTH = 4096
PERF_FILE = "./examples/models/performance.csv"
PERF_COLUMNS = "batch, input_len, output_len, embedding_len, first_token_time(ms), \
non_first_token_time(ms), ResponseTime(ms),E2E Throughput Average(Tokens/s)\n"
PRED_FILE = "./examples/models/result.json"
MICROBATCH_SIZE = "microbatch_size"
class PARunner:
def __init__(self, **kwargs):
from datetime import timedelta
dist.init_process_group(
backend="hccl",
world_size=kwargs.get('world_size', '1'),
rank=kwargs.get('rank', '0'),
timeout=timedelta(seconds=1800)
)
logger.info(f"dgi: kwargs:{kwargs}")
self.rank = kwargs.get('rank', '0')
logger.info(f"Initialized HCCL process group for rank {self.rank}")
self.local_rank = kwargs.get('local_rank', self.rank)
self.world_size = kwargs.get('world_size', '1')
self.model_path = kwargs.get('model_path', None)
self.input_text = kwargs.get('input_text', None)
self.max_prefill_batch_size = kwargs.get('max_prefill_batch_size', None)
self.max_batch_size = kwargs.get('max_batch_size', None)
self.max_input_length = kwargs.get('max_input_length', None)
self.max_output_length = kwargs.get('max_output_length', None)
self.max_position_embeddings = kwargs.get('max_position_embeddings', None)
self.max_prefill_tokens = kwargs.get('max_prefill_tokens', None)
self.enable_atb_torch = kwargs.get('enable_atb_torch', False)
self.trust_remote_code = kwargs.get('trust_remote_code', False)
self.block_size = kwargs.get('block_size', None)
self.chat_template = kwargs.get('chat_template', None)
self.load_tokenizer = kwargs.get('load_tokenizer', True)
self.prefill_length = kwargs.get('prefill_length', 8192)
self.distributed_enable = kwargs.get('distributed_enable', False)
self.llm_config_path = kwargs.get('llm_config_path', None)
self.check_limits()
kw_args = kwargs.get('kw_args', '')
if len(kw_args) > MAX_KW_ARGS_LENGTH:
raise ValueError(f'The length of kw_args is too long and should be within (0, {MAX_KW_ARGS_LENGTH}]')
if kw_args:
try:
kw_args = json.loads(kw_args)
except json.decoder.JSONDecodeError:
print_log(self.local_rank, logger.warning, 'Please checkout input of kw_args!')
kw_args = {}
else:
kw_args = {}
if kwargs.get(MICROBATCH_SIZE, -1) == -1:
microbatch_size = self.max_batch_size
elif kwargs.get(MICROBATCH_SIZE, -1) <= self.max_batch_size and kwargs.get(MICROBATCH_SIZE, -1) > 0:
microbatch_size = kwargs.get(MICROBATCH_SIZE, -1)
else:
raise ValueError(f'Micro batch size should be within [1, {self.max_batch_size}] or -1!')
if 'num_speculative_tokens' not in kw_args:
kw_args['num_speculative_tokens'] = int(ENV.deepseek_mtp)
self.model = ModelRunner(
self.model_path, rank=self.rank, world_size=self.world_size,
local_rank=self.local_rank,
max_position_embeddings=self.max_position_embeddings,
enable_atb_torch=self.enable_atb_torch,
load_tokenizer=self.load_tokenizer,
lora_modules=kwargs.get('lora_modules', None),
max_loras=kwargs.get('max_loras', 0),
max_lora_rank=kwargs.get('max_lora_rank', 0),
trust_remote_code=self.trust_remote_code,
llm_config_path=self.llm_config_path,
dp=kwargs.get("dp", -1),
tp=kwargs.get("tp", -1),
moe_tp=kwargs.get("moe_tp", -1),
pp=kwargs.get("pp", -1),
microbatch_size=microbatch_size,
moe_ep=kwargs.get("moe_ep", -1),
sp=kwargs.get("sp", -1),
cp=kwargs.get("cp", -1),
distributed_enable=self.distributed_enable,
max_batch_size=self.max_batch_size,
**kw_args
)
self.tokenizer = self.model.tokenizer
if self.chat_template:
self.tokenizer.chat_template = self._load_chat_template(self.chat_template)
self.dtype = self.model.dtype
self.quantize = self.model.quantize
self.kv_quant_type = self.model.kv_quant_type
self.fa_quant_type = self.model.fa_quant_type
self.model.load_weights(**kw_args)
self.device = self.model.device
self.model_config = ModelConfig(self.model.num_heads,
self.model.num_kv_heads,
self.model.config.num_key_value_heads \
if hasattr(self.model.config, 'num_key_value_heads') \
else self.model.num_kv_heads,
self.model.k_head_size,
self.model.v_head_size,
self.model.num_layers,
self.model.device,
self.model.dtype,
self.model.soc_info,
self.kv_quant_type,
self.fa_quant_type,
self.model.mapping,
self.model.kvcache_quant_layers,
getattr(self.model.config, 'cla_share_factor', 1),
getattr(self.model.config, 'model_type', None),
self.model.enable_nz)
self.max_memory = NpuHbmInfo.get_hbm_capacity(self.local_rank, self.world_size, self.model.soc_info.need_nz)
self.init_memory = int(
self.max_memory * NpuHbmInfo.get_hbm_usage(self.local_rank, self.world_size, self.model.soc_info.need_nz))
print_log(self.rank, logger.info, f'hbm_capacity(GB): {self.max_memory / (1024 ** 3)}, '
f'init_memory(GB): {self.init_memory / (1024 ** 3)}')
self.warm_up_memory = 0
self.warm_up_num_blocks = 0
self.cache_manager = None
self.hidden_state_cache = None
self.compress_head_enable = ENV.compress_head_enable
self.omni_attention_enable = ENV.omni_attention_enable
self.eplb_forwarder = None
eplb_level = getattr(self.model.model, "eplb_level", 0)
if eplb_level == EPLBType.DYNAMIC_EPLB:
self.eplb_worker = EplbWorker(self.model, self.rank, self.model_path, self.device)
self.eplb_forwarder = self.eplb_worker.eplb_forwarder
def test_broadcast_data(self):
if self.rank != 0:
data = [None]
else:
data = [1]
logger.info("before broadcast")
dist.broadcast_object_list(data, src=0)
logger.info("after broadcast")
def __repr__(self):
return (
"PARunner("
+ f"model_path={self.model_path}, "
+ f"input_text={self.input_text}, "
+ f"max_position_embeddings={self.max_position_embeddings}, "
+ f"max_input_length={self.max_input_length}, "
+ f"max_output_length={self.max_output_length}, "
+ f"max_prefill_tokens={self.max_prefill_tokens}, "
+ f"load_tokenizer={self.load_tokenizer}, "
+ f"enable_atb_torch={self.enable_atb_torch}, "
+ f"max_prefill_batch_size={self.max_prefill_batch_size}, "
+ f"max_batch_size={self.max_batch_size}, "
+ f"dtype={self.dtype}, "
+ f"block_size={self.block_size}, "
+ f"model_config={self.model_config}, "
+ f"max_memory={self.max_memory}, "
)
@staticmethod
def save_input_output_ids(req_list):
for idx, req in enumerate(req_list):
input_ids_save_filename = f"input_ids_{idx}.pth"
input_ids_save_path = os.path.join(ENV.token_ids_save_folder, input_ids_save_filename)
input_ids_save_path = file_utils.standardize_path(input_ids_save_path)
file_utils.check_file_safety(input_ids_save_path, 'w', is_check_file_size=False)
torch.save(req.input_ids.cpu(), input_ids_save_path)
os.chmod(input_ids_save_path, 0o640)
output_ids_save_filename = f"output_ids_{idx}.txt"
output_ids_save_path = os.path.join(ENV.token_ids_save_folder, output_ids_save_filename)
with file_utils.safe_open(output_ids_save_path, 'w', encoding='utf-8') as f:
f.write(' '.join(map(str, req.out_token_list)))
@staticmethod
def _update_inputs_items(inputs):
input_texts, input_ids, input_conversations, input_adapters = [], [], [], []
if isinstance(inputs, list) and inputs:
if isinstance(inputs[0], str):
input_texts = inputs
elif isinstance(inputs[0], torch.Tensor):
input_ids = inputs
elif isinstance(inputs[0], list) and inputs[0]:
if isinstance(inputs[0][0], int):
input_ids = inputs
elif isinstance(inputs[0][0], dict):
input_conversations = inputs
elif isinstance(inputs[0], dict):
prompt_key = "prompt"
if isinstance(inputs[0].get(prompt_key), torch.Tensor):
input_ids = [input.get(prompt_key) for input in inputs]
else:
input_texts = [input.get(prompt_key) for input in inputs]
input_adapters = [input.get("adapter") for input in inputs]
return input_texts, input_ids, input_conversations, input_adapters
@staticmethod
def _load_chat_template(chat_template: str):
if os.path.exists(chat_template):
with file_utils.safe_open(chat_template, 'r', encoding='utf-8') as f:
chat_template_content = f.read()
else:
chat_template_content = chat_template
return chat_template_content
def check_limits(self):
if self.max_input_length <= 0:
print_log(self.rank, logger.warning,
f"'max_input_length' must be greater than 0, but is {self.max_input_length}. "
"Please increase the value of 'max_input_length'. ")
if self.max_output_length <= 0:
print_log(self.rank, logger.warning,
f"'max_output_length' must be greater than 0, but is {self.max_output_length}. "
"Please increase the value of 'max_output_length'. ")
if self.max_position_embeddings and \
self.max_position_embeddings < self.max_input_length + self.max_output_length:
print_log(self.rank, logger.warning,
f"'max_position_embeddings' must be no less than "
f"{self.max_input_length + self.max_output_length}, but is {self.max_position_embeddings}. "
"Please increase the value of 'max_position_embeddings'. ")
if self.max_prefill_batch_size and self.max_prefill_batch_size <= 0:
print_log(self.rank, logger.warning,
f"'max_prefill_batch_size' must be greater than 0, but is {self.max_prefill_batch_size}. "
"Please increase the value of 'max_prefill_batch_size'. ")
if self.max_batch_size <= 0:
print_log(self.rank, logger.warning,
f"'max_batch_size' must be greater than 0, but is {self.max_batch_size}. "
"Please increase the value of 'max_batch_size'. ")
if self.block_size > OPERATOR_BOUND:
print_log(self.rank, logger.warning,
f"'block_size' must be lower than or equal to 576, but is {self.block_size}. "
"Please decrease the value of 'block_size'. ")
if self.block_size <= 0:
print_log(self.rank, logger.warning,
f"'block_size' is too small: {self.block_size}. "
"Please increase the value of 'block_size'. ")
if self.max_prefill_tokens != -1 and \
self.max_prefill_tokens < self.max_input_length:
print_log(self.rank, logger.warning,
f"'max_prefill_tokens' must be greater than or equal to {self.max_input_length}, "
f"but is {self.max_prefill_tokens}. Please increse the value of 'max_prefill_tokens'. ")
def calc_max_length(self):
if self.model.mapping.has_attn_cp() and self.model.mapping.has_attn_inner_sp():
max_length = math.ceil((math.ceil(self.max_input_length / self.model.mapping.attn_cp.group_size) +
self.max_output_length) / self.model.mapping.attn_inner_sp.group_size)
elif self.model.mapping.has_attn_inner_sp():
max_length = math.ceil((self.max_input_length + self.max_output_length) /
self.model.mapping.attn_inner_sp.group_size)
elif self.model.mapping.has_attn_cp():
# CP will perform padding on the input.
max_length = (math.ceil(self.max_input_length / self.model.mapping.attn_cp.group_size) +
math.ceil(self.max_output_length / self.model.mapping.attn_cp.group_size))
else:
max_length = self.max_input_length + self.max_output_length
return max_length
def warm_up(self):
self.model.model.warmup_is_end = False
if self.max_prefill_batch_size is None:
self.max_prefill_batch_size = self.max_batch_size
if ENV.modeltest_pd_split_enable:
if self.max_input_length > self.prefill_length:
self.max_prefill_batch_size = 1
else:
self.max_prefill_batch_size = math.floor(self.prefill_length / self.max_input_length)
if self.max_prefill_batch_size > self.max_batch_size:
self.max_prefill_batch_size = self.max_batch_size
effective_max_prefill_batch_size = math.ceil(
self.max_prefill_batch_size / self.model.mapping.attn_dp.group_size)
effective_max_batch_size = math.ceil(
self.max_batch_size / self.model.mapping.attn_dp.group_size)
if self.max_prefill_tokens == -1:
self.max_prefill_tokens = effective_max_prefill_batch_size * (self.max_input_length +
self.max_output_length)
print_log(self.rank, logger.info, "---------------begin warm_up---------------")
try:
max_length = self.calc_max_length()
self.warm_up_num_blocks = math.ceil(max_length / self.block_size) * effective_max_batch_size
except ZeroDivisionError as e:
raise ZeroDivisionError from e
cache_config = CacheConfig(self.warm_up_num_blocks, self.block_size)
if self.compress_head_enable or self.omni_attention_enable:
cache_config = CacheConfig(self.warm_up_num_blocks, self.block_size, \
self.max_input_length, self.max_output_length, effective_max_batch_size, self.rank, self.world_size)
self.cache_manager = CacheManager(cache_config, self.model_config)
if ENV.deepseek_mtp:
batch_size_for_cache = max(self.max_batch_size, self.max_prefill_batch_size)
self.cache_manager.hidden_state_cache = HiddenStateCache(self.model.config.hidden_size,
ENV.deepseek_mtp, batch_size_for_cache, self.model.dtype, self.cache_manager.device)
else:
self.hidden_state_cache = None
self.model.postprocessor.max_new_tokens = 1
input_ids = torch.ones(self.max_input_length, dtype=torch.int64)
req_list = [
request_from_token(input_ids, 1, self.block_size, req_idx=idx, mapping=self.model.mapping)
for idx in range(self.max_prefill_batch_size)
]
generate_req(req_list, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager,
self.eplb_forwarder)
self.warm_up_memory = int(
self.max_memory * NpuHbmInfo.get_hbm_usage(self.local_rank, self.world_size, self.model.soc_info.need_nz))
print_log(self.rank, logger.info, f'warmup_memory(GB): {self.warm_up_memory / (1024 ** 3): .2f}')
print_log(self.rank, logger.info, "---------------end warm_up---------------")
self.model.model.warmup_is_end = True
def infer(self, inputs, batch_size, max_output_length, ignore_eos, is_chat_model=False, **kwargs):
if torch.distributed.is_initialized():
logger.info(">>>> torch.distributed is already initialized.")
else:
logger.error(">>>> torch.distributed has not been initialized!")
print_log(self.rank, logger.info, "---------------begin inference---------------")
if ignore_eos:
self.model.postprocessor.eos_token_id = []
is_truncation = kwargs.get("truncation", False)
input_ids, adapter_ids = self._build_model_inputs(inputs, is_chat_model, is_truncation)
if len(input_ids) == 1:
req_list = [request_from_token(input_ids[0], max_output_length, self.block_size, req_idx=idx,
adapter_id=adapter_ids[0], mapping=self.model.mapping,
postprocessor=self.model.postprocessor)
for idx in range(batch_size)]
else:
req_list = [request_from_token(input_ids_ins, max_output_length, self.block_size, req_idx=idx,
adapter_id=adapter_id, mapping=self.model.mapping, postprocessor=self.model.postprocessor)
for idx, (input_ids_ins, adapter_id) in enumerate(zip(input_ids, adapter_ids))]
print_log(self.rank, logger.debug, f'req_list[0].input_ids: {req_list[0].input_ids}')
if not self.cache_manager:
if self.max_prefill_tokens == -1:
self.max_prefill_tokens = self.max_batch_size * (self.max_input_length + self.max_output_length)
cache_block_size = self.block_size * self.model.num_kv_heads * self.model.head_size
dtype_size = CacheManager.get_dtype_size(self.dtype)
total_cache_size = self.model.num_layers * cache_block_size * 2 * dtype_size
max_memory = ENV.memory_fraction * self.max_memory
free_memory = max_memory - ENV.reserved_memory_gb * (1 << 30) - (
self.warm_up_memory if self.warm_up_memory != 0 else self.init_memory)
print_log(self.rank, logger.info,
f"infer max_memory(GB): {max_memory / (1024 ** 3): .2f}, "
f"warm_up_memory(GB): {self.warm_up_memory / (1024 ** 3): .2f}, "
f"free_memory(GB): {free_memory / (1024 ** 3): .2f}")
try:
num_blocks = int(free_memory // total_cache_size)
except ZeroDivisionError as e:
raise ZeroDivisionError("total_cache_size should not be 0") from e
print_log(self.rank, logger.info, f"num_blocks: {num_blocks}, free_memory: {free_memory}")
cache_config = CacheConfig(num_blocks, self.block_size)
if self.compress_head_enable:
cache_config = CacheConfig(self.warm_up_num_blocks, self.block_size, \
self.max_input_length, self.max_output_length, self.max_batch_size, self.rank, self.world_size)
self.cache_manager = CacheManager(cache_config, self.model_config)
if ENV.benchmark_enable:
req_list_dummy = copy.deepcopy(req_list)
self.model.postprocessor.max_new_tokens = 2
generate_req(req_list_dummy, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager,
self.eplb_forwarder)
self.model.postprocessor.max_new_tokens = max_output_length
skip_special_tokens = kwargs.get("skip_special_tokens", False)
if not ENV.profiling_enable:
print_log(self.rank, logger.debug, "No profiling")
torch.npu.synchronize()
e2e_start = time.time()
generate_req(req_list, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager,
self.eplb_forwarder)
_, _ = decode_token(req_list, self.tokenizer, skip_special_tokens)
torch.npu.synchronize()
e2e_end = time.time()
e2e_time = e2e_end - e2e_start
else:
print_log(self.rank, logger.debug, "Enter profiling")
profiling_path = ENV.profiling_filepath
profiler_level = torch_npu.profiler.ProfilerLevel
actual_profiler_level = getattr(profiler_level, ENV.profiling_level)
torch.npu.synchronize()
e2e_start = time.time()
experimental_config = torch_npu.profiler._ExperimentalConfig(
aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization,
profiler_level=actual_profiler_level,
l2_cache=False,
data_simplification=False
)
with torch_npu.profiler.profile(
activities=[
torch_npu.profiler.ProfilerActivity.CPU,
torch_npu.profiler.ProfilerActivity.NPU
],
on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profiling_path),
record_shapes=True,
profile_memory=True,
with_stack=False,
with_flops=False,
with_modules=False,
experimental_config=experimental_config):
generate_req(req_list, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager,
self.eplb_forwarder)
torch.npu.synchronize()
e2e_end = time.time()
e2e_time = e2e_end - e2e_start
generate_text_list, token_num_list = decode_token(req_list, self.tokenizer, skip_special_tokens)
if ENV.modeltest_dataset_specified:
ENV.update()
if ENV.token_ids_save_enable:
if self.model.mapping.has_pp():
if self.rank == self.world_size - 1:
self.save_input_output_ids(req_list)
else:
if self.local_rank == 0:
self.save_input_output_ids(req_list)
print_log(self.rank, logger.info, "---------------end inference---------------")
return generate_text_list, token_num_list, e2e_time
def shutdown(self):
eplb_level = getattr(self.model.config, "eplb_level", 0)
if eplb_level == EPLBType.DYNAMIC_EPLB:
self.eplb_worker.eplb_loader.shutdown()
self.eplb_worker.eplb_planner.shutdown()
def _build_model_inputs(self, inputs, is_chat_model, is_truncation=False):
input_texts, input_ids, input_conversations, input_adapters = self._update_inputs_items(inputs)
if not (input_texts or input_ids or input_conversations):
raise ValueError(f"The inputs of `PARunner.infer` must be as List[str], List[torch.Tensor], List[List[int]]"
f" or List[List[Dict]]. Now the inputs ({inputs}) is not acceptable or is empty.")
if not input_ids and not self.tokenizer:
raise ValueError("Can‘t encode input_texts, because load_tokenizer=False, "
"please enter input_ids or enable load_tokenizer")
if is_chat_model:
if input_conversations:
input_ids = self.model.build_inputs(input_conversations)
elif input_texts:
input_conversations = [[{"role": "user", "content": t}] for t in input_texts]
input_ids = self.model.build_inputs(input_conversations)
else:
print_log(self.rank, logger.warning, "Neither conversations nor input_texts exist, "
"'chat' parameter is not effective.")
elif input_texts:
input_ids = [self.tokenizer([text], return_tensors="pt", truncation=is_truncation)["input_ids"].flatten()
for text in input_texts]
# token长度校验
for item in input_ids:
if len(item) > self.max_input_length:
print_log(self.rank, logger.warning,
"Num of tokens in input_text is larger than max_input_length. "
"Please shorten input_text to avoid out of memory.")
if not input_adapters:
input_adapters = [None] * len(input_ids)
return input_ids, input_adapters
def cmd_bool(cmd_arg):
if cmd_arg == "True":
return True
elif cmd_arg == "False":
return False
raise ValueError(f"{cmd_arg} should be a boolean")
def parse_list_of_json(list_json):
return json.loads(list_json)
def parse_ids(list_str):
return [int(item) for item in list_str.split(',')]
def input_texts_parser(value):
if os.path.isfile(value):
with file_utils.safe_open(value, 'r') as opened_file:
return opened_file.read()
else:
return value
#define Argument Parser
def parse_arguments():
store_true = 'store_true'
parser = ArgumentParser()
parser.add_argument('--model_path', type=str, help="model and tokenizer path")
parser.add_argument(
'--input_texts',
type=input_texts_parser,
nargs='+',
default=["What's deep learning?"])
parser.add_argument(
'--input_ids',
type=parse_ids,
nargs='+',
default=None)
parser.add_argument(
'--input_file',
type=str,
help='This parameter is used to input multi-turn dialogue information in the form '
'of a jsonl file, with each line in the format of a List[Dict]. Each dictionary '
'(Dict) must contain at least two fields: "role" and "content".',
default=None)
parser.add_argument(
'--input_dict',
help="Lora input, accepted format: "
"'[{\"prompt\": \"prompt in text\", \"adapter\": \"adapter id defined in lora_adapter param\"}]'",
type=parse_list_of_json,
default=None)
parser.add_argument("--max_prefill_batch_size", type=int, default=None)
parser.add_argument("--max_batch_size", type=int, default=1)
parser.add_argument('--max_input_length', type=int, default=1024)
parser.add_argument('--max_output_length', type=int, default=20)
parser.add_argument('--max_position_embeddings', type=int, default=None)
parser.add_argument('--max_prefill_tokens', type=int, default=-1)
parser.add_argument("--block_size", type=int, default=128)
parser.add_argument('--chat_template', type=str, default=None)
parser.add_argument('--ignore_eos', action=store_true)
parser.add_argument('--is_chat_model', action=store_true)
parser.add_argument('--is_embedding_model', action=store_true)
parser.add_argument('--load_tokenizer', type=cmd_bool, default=True)
parser.add_argument('--enable_atb_torch', action=store_true)
parser.add_argument('--trust_remote_code', action=store_true)
parser.add_argument('--dp', type=int, default=-1)
parser.add_argument('--tp', type=int, default=-1)
parser.add_argument('--moe_tp', type=int, default=-1)
parser.add_argument('--pp', type=int, default=-1)
parser.add_argument('--microbatch_size', type=int, default=-1)
parser.add_argument('--moe_ep', type=int, default=-1)
parser.add_argument('--sp', type=int, default=-1)
parser.add_argument('--cp', type=int, default=-1)
parser.add_argument('--kw_args', type=str, default='', help='json input')
parser.add_argument('--lora_modules', type=str, default=None, help="Define lora modules, accepted format:"
"'{\"adapter1\": \"/path/to/lora1\", \"adapter2\": \"/path/to/lora2\"}'",
validator=StringArgumentValidator(min_length=2, max_length=MAX_JSON_LENGTH, allow_none=True))
parser.add_argument('--max_loras', type=int, default=0, help="Define the MAX number of loRAs to store in NPU memory",
validator=NumberArgumentValidator(min_value=0))
parser.add_argument('--max_lora_rank', type=int, default=0, help="Define the MAX loRA rank",
validator=NumberArgumentValidator(min_value=0))
parser.add_argument('--distributed_enable', action=store_true)
return parser.parse_args()
def test_infer(pa_runner):
infer_params = {
"inputs": ["What's deep learning?"],
"batch_size": 1,
"max_output_length": 10,
"ignore_eos": True,
"is_chat_model": True
}
generate_texts, token_nums, _ = pa_runner.infer(**infer_params)
length = len(infer_inputs)
for i, generate_text in enumerate(generate_texts):
if input_dict['is_embedding_model']:
embedding_tensor_path = f"{os.getcwd()}/examples/embedding_tensor"
print_log(rank, logger.info, f"Context[{i}]: \nembedding tensor path is: {embedding_tensor_path}")
else:
print_log(rank, logger.info, f'Answer[{i}]: {generate_text}')
print_log(rank, logger.info, f'Generate[{i}] token num: {token_nums[i]}')
if __name__ == '__main__':
args = parse_arguments()
rank = ENV.rank
local_rank = ENV.local_rank
world_size = ENV.world_size
input_dict = {
'rank': rank,
'world_size': world_size,
'local_rank': local_rank,
**vars(args)
}
# 输入优先级input_dict > input_ids > input_texts
if args.input_dict:
infer_inputs = args.input_dict
elif args.input_ids:
infer_inputs = args.input_ids
else:
infer_inputs = args.input_texts
if args.is_chat_model and args.input_file:
conversations = []
with file_utils.safe_open(args.input_file, 'r', encoding='utf-8') as file:
for line in file_utils.safe_readlines(file):
data_line = json.loads(line)
conversations.append(data_line)
infer_inputs = conversations
pa_runner = PARunner(**input_dict)
print_log(rank, logger.info, f'pa_runner: {pa_runner}')
# pa_runner.test_broadcast_data()
print_log(rank, logger.info, f'before_test_infer')
test_infer(pa_runner)
test_infer(pa_runner)
print_log(rank, logger.info, f'after_test_infer')
print_log(rank, logger.info, f'before_test_infer1')
test_infer(pa_runner)
print_log(rank, logger.info, f'after_test_infer1')
# exit(0)
pa_runner.warm_up()
infer_params = {
"inputs": infer_inputs,
"batch_size": args.max_batch_size,
"max_output_length": args.max_output_length,
"ignore_eos": args.ignore_eos,
"is_chat_model": args.is_chat_model
}
generate_texts, token_nums, _ = pa_runner.infer(**infer_params)
length = len(infer_inputs)
for i, generate_text in enumerate(generate_texts):
if input_dict['is_embedding_model']:
embedding_tensor_path = f"{os.getcwd()}/examples/embedding_tensor"
print_log(rank, logger.info, f"Context[{i}]: \nembedding tensor path is: {embedding_tensor_path}")
else:
print_log(rank, logger.info, f'Answer[{i}]: {generate_text}')
print_log(rank, logger.info, f'Generate[{i}] token num: {token_nums[i]}')
# 下面的命令两台主机都要执行,注意:1、run_pa的挂载(容器内修改文件也可)2.rank_table_file的挂载
docker run -it \
--name dgi-dev-mindie \
--network=host \
--privileged=true \
--device=/dev/davinci_manager \
--device=/dev/hisi_hdc \
--device=/dev/devmm_svm \
-v /usr/local/lib/node_modules/:/usr/local/lib/node_modules/:ro \
-v /usr/local/bin/node:/usr/local/bin/node:ro \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64/common \
-v /usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/driver \
-v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v /etc/vnpu.cfg:/etc/vnpu.cfg \
-v /usr/local/sbin:/usr/local/sbin \
-v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \
-v /data1/:/data1/ \
-v /data2/:/data2/ \
-v /data1/MindIE_Auto_Deploy_Tool/rank_table_file.json:/data1/MindIE_Auto_Deploy_Tool/rank_table_file.json \
-v /data1/atb-models/examples/run_pa.py:/usr/local/Ascend/atb-models/examples/run_pa.py \
-w /root \
-u root \
--shm-size=64GB \
--ulimit msgqueue=-1 \
-e ASCEND_RT_VISIBLE_DEVICES="0,1,2,3,4,5,6,7" \
swr.cn-south-1.myhuaweicloud.com/ascendhub/mindie:2.2.RC1-800I-A2-py311-openeuler24.03-lts
# 进容器后执行下面的命令
source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh
source /usr/local/Ascend/atb-models/set_env.sh
export ATB_LLM_HCCL_ENABLE=1
export ATB_LLM_COMM_BACKEND="hccl"
export HCCL_CONNECT_TIMEOUT=7200
export HCCL_EXEC_TIMEOUT=0
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export ATB_LLM_BENCHMARK_ENABLE=1
export ATB_LLM_ENABLE_AUTO_TRANSPOSE=0
# 打开 ATB 日志
# export ASCEND_GLOBAL_LOG_LEVEL=3
# export ASCEND_SLOG_PRINT_TO_STDOUT=1
# MINDIE日志
export MINDIE_LOG_TO_STDOUT=1
export HCCL_DETERMINISTIC=false
export LCCL_DETERMINISTIC=0
export HCCL_BUFFSIZE=120
export ATB_WORKSPACE_MEM_ALLOC_GLOBAL=1
# 使用ranktable!
export RANK_TABLE_FILE=/data1/MindIE_Auto_Deploy_Tool/rank_table_file.json
首先是使用rank_table推理,不用broadcast
主机执行:
# 主机执行
torchrun \
--nproc_per_node 8 \
--master_port 12345 \
--nnodes 2 \
--node_rank 0 \
--master_addr 10.17.44.64 \
/usr/local/Ascend/atb-models/examples/run_pa.py \
--model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/master.log
从机执行:
# 从机执行
torchrun \
--nproc_per_node 8 \
--master_port 12345 \
--nnodes 2 \
--node_rank 1 \
--master_addr 10.17.44.64 \
/usr/local/Ascend/atb-models/examples/run_pa.py \
--model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/slave.log
2026-01-25 19:55启动
日志中可能报警告:
[2026-01-25 19:55:53.460882] [833] [833] [llmmodels] [WARN] [context_factory.cpp:37] Failed to aclInit: 100002
但结果正常推理,输出
[2026-01-25 19:56:42,324] [829] [281473611860544] [llmmodels] [INFO] [run_pa.py-652] : Answer[0]: What's the difference between deep learning and machine learning? What's the difference between deep learning and neural
[2026-01-25 19:56:42,324] [829] [281473611860544] [llmmodels] [INFO] [run_pa.py-653] : Generate[0] token num: (0, 20)
使用rank_table推理,推理前加入broadcast
两台容器同时进vim /usr/local/Ascend/atb-models/examples/run_pa.py +626
取消注释pa_runner.test_broadcast_data()
主机执行:
# 主机执行
torchrun \
--nproc_per_node 8 \
--master_port 12345 \
--nnodes 2 \
--node_rank 0 \
--master_addr 10.17.44.64 \
/usr/local/Ascend/atb-models/examples/run_pa.py \
--model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/master1.log
从机执行:
# 从机执行
torchrun \
--nproc_per_node 8 \
--master_port 12345 \
--nnodes 2 \
--node_rank 1 \
--master_addr 10.17.44.64 \
/usr/local/Ascend/atb-models/examples/run_pa.py \
--model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/slave1.log
2026-01-25 20:01启动,报错:
[2026-01-25 20:02:05,455] [3754] [281473078336064] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory
/usr/local/lib64/python3.11/site-packages/torch_npu/utils/storage.py:41: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly. To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
if self.device.type != 'cpu':
[2026-01-25 20:02:05,521] [3754] [281473078336064] [llmmodels] [INFO] [generate.py-1130] : ------total req num: 1, infer start--------
[2026-01-25 20:02:05,521] [3758] [281473476311616] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:02:06,782] [3755] [281473268972096] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:02:07,495] [3760] [281473129114176] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:02:08,468] [3761] [281472942824000] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:02:08,474] [3757] [281473277147712] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:02:10,166] [3759] [281473160243776] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:02:13,227] [3756] [281473215760960] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:02:15,449] [3759] [281473160243776] [llmmodels] [INFO] [run_pa.py-165] : after broadcast
[2026-01-25 20:02:15,449] [3755] [281473268972096] [llmmodels] [INFO] [run_pa.py-165] : after broadcast
[2026-01-25 20:02:15,449] [3758] [281473476311616] [llmmodels] [INFO] [run_pa.py-165] : after broadcast
[2026-01-25 20:02:15,450] [3759] [281473160243776] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized.
[2026-01-25 20:02:15,450] [3760] [281473129114176] [llmmodels] [INFO] [run_pa.py-165] : after broadcast
[2026-01-25 20:02:15,450] [3755] [281473268972096] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized.
[2026-01-25 20:02:15,450] [3758] [281473476311616] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized.
[2026-01-25 20:02:15,450] [3760] [281473129114176] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized.
[2026-01-25 20:02:15,450] [3757] [281473277147712] [llmmodels] [INFO] [run_pa.py-165] : after broadcast
[2026-01-25 20:02:15,450] [3757] [281473277147712] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized.
[2026-01-25 20:02:15,454] [3756] [281473215760960] [llmmodels] [INFO] [run_pa.py-165] : after broadcast
[2026-01-25 20:02:15,454] [3756] [281473215760960] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized.
[2026-01-25 20:02:15,455] [3761] [281472942824000] [llmmodels] [INFO] [run_pa.py-165] : after broadcast
[2026-01-25 20:02:15,455] [3761] [281472942824000] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized.
[2026-01-25 20:02:15,553] [3759] [281473160243776] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory
[2026-01-25 20:02:15,555] [3755] [281473268972096] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory
[2026-01-25 20:02:15,556] [3761] [281472942824000] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory
[2026-01-25 20:02:15,562] [3758] [281473476311616] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory
[2026-01-25 20:02:15,564] [3756] [281473215760960] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory
[2026-01-25 20:02:15,565] [3757] [281473277147712] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory
[2026-01-25 20:02:15,567] [3760] [281473129114176] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory
mki_log mkdir /root/ascend/log/atb
Traceback (most recent call last):
File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 628, in <module>
test_infer(pa_runner)
File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 584, in test_infer
generate_texts, token_nums, _ = pa_runner.infer(**infer_params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 390, in infer
generate_req(req_list_dummy, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager,
File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 1218, in generate_req
generate_token_with_clocking(model, cache_manager, batch, eplb_forwarder)
File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 892, in generate_token_with_clocking
res = generate_token(model, cache_manager, input_batch_in, eplb_forwarder)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 616, in generate_token
logits = model.forward(
^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/atb_llm/runner/model_runner.py", line 324, in forward
res = self.model.forward(**kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/atb_llm/models/qwen2_moe/flash_causal_qwen2_moe.py", line 438, in forward
logits = super().forward(**kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/atb_llm/models/base/flash_causal_lm.py", line 519, in forward
self.init_ascend_weight()
File "/usr/local/Ascend/atb-models/atb_llm/models/qwen2_moe/flash_causal_qwen2_moe.py", line 311, in init_ascend_weight
self.acl_encoder_operation.set_param(self.acl_param_encoder)
RuntimeError: External Comm Manager: Create the hccl communication group failed. export ASCEND_GLOBAL_LOG_LEVEL=3, export ASCEND_SLOG_PRINT_TO_STDOUT=1 to see more details. Default log path is $HOME/atb/log.
Traceback (most recent call last):
File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 628, in <module>
test_infer(pa_runner)
File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 584, in test_infer
generate_texts, token_nums, _ = pa_runner.infer(**infer_params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 390, in infer
generate_req(req_list_dummy, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager,
File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 1218, in generate_req
generate_token_with_clocking(model, cache_manager, batch, eplb_forwarder)
File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 892, in generate_token_with_clocking
res = generate_token(model, cache_manager, input_batch_in, eplb_forwarder)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 616, in generate_token
logits = model.forward(
^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/atb_llm/runner/model_runner.py", line 324, in forward
res = self.model.forward(**kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/atb_llm/models/qwen2_moe/flash_causal_qwen2_moe.py", line 438, in forward
logits = super().forward(**kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/atb_llm/models/base/flash_causal_lm.py", line 519, in forward
self.init_ascend_weight()
File "/usr/local/Ascend/atb-models/atb_llm/models/qwen2_moe/flash_causal_qwen2_moe.py", line 311, in init_ascend_weight
self.acl_encoder_operation.set_param(self.acl_param_encoder)
RuntimeError: External Comm Manager: Create the hccl communication group failed. export ASCEND_GLOBAL_LOG_LEVEL=3, export ASCEND_SLOG_PRINT_TO_STDOUT=1 to see more details. Default log path is $HOME/atb/log.
使用rank_table推理,推理后尝试broadcast:
两台容器同时进vim /usr/local/Ascend/atb-models/examples/run_pa.py +626
将626行的pa_runner.test_broadcast_data()移动628行到两个test_infer之间,像是这样:
624 pa_runner = PARunner(**input_dict)
625 print_log(rank, logger.info, f'pa_runner: {pa_runner}')
626 print_log(rank, logger.info, f'before_test_infer')
627 test_infer(pa_runner)
628 pa_runner.test_broadcast_data()
629 test_infer(pa_runner)
630 print_log(rank, logger.info, f'after_test_infer')
主机执行:
# 主机执行
torchrun \
--nproc_per_node 8 \
--master_port 12345 \
--nnodes 2 \
--node_rank 0 \
--master_addr 10.17.44.64 \
/usr/local/Ascend/atb-models/examples/run_pa.py \
--model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/master2.log
从机执行:
# 从机执行
torchrun \
--nproc_per_node 8 \
--master_port 12345 \
--nnodes 2 \
--node_rank 1 \
--master_addr 10.17.44.64 \
/usr/local/Ascend/atb-models/examples/run_pa.py \
--model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/slave2.log
2026-01-25 20:08 启动,第一个test_infer正常,程序卡死在before_broadcast与after_broadcast之间,日志例如(从机,主机同样):等了十分钟还是没响应,ctrl+c停掉了。
[2026-01-25 20:09:56,791] [5677] [281473705536064] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:09:56,792] [5676] [281472899271232] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:09:56,792] [5678] [281473714686528] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:09:56,792] [5672] [281473198320192] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:09:56,792] [5674] [281472934701632] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:09:56,792] [5675] [281473338534464] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:09:56,792] [5673] [281473130969664] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
[2026-01-25 20:09:56,792] [5679] [281473051490880] [llmmodels] [INFO] [run_pa.py-163] : before broadcast
在上一个的基础上不使用rank_table推理
紧接上一个例子,两个test_infer之间使用broadcast,并且unset RANK_TABLE_FILE变量
unset RANK_TABLE_FILE
主机执行:
# 主机执行
torchrun \
--nproc_per_node 8 \
--master_port 12345 \
--nnodes 2 \
--node_rank 0 \
--master_addr 10.17.44.64 \
/usr/local/Ascend/atb-models/examples/run_pa.py \
--model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/master3.log
从机执行:
# 从机执行
torchrun \
--nproc_per_node 8 \
--master_port 12345 \
--nnodes 2 \
--node_rank 1 \
--master_addr 10.17.44.64 \
/usr/local/Ascend/atb-models/examples/run_pa.py \
--model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/slave3.log
2026-01-25 20:20启动,报错变成了:
[2026-01-25 20:21:21,470] [7981] [281473698683456] [llmmodels] [INFO] [generate.py-1130] : ------total req num: 1, infer start--------
[2026-01-25 20:21:21,969] [7985] [281473459296832] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized.
[2026-01-25 20:21:22,102] [7985] [281473459296832] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory
[2026-01-25 20:21:22.214388] [7981] [7981] [llmmodels] [INFO] [moe_decoder_model.cpp:211] model enableSplitFuse: 0
Traceback (most recent call last):
File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 627, in <module>
test_infer(pa_runner)
File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 584, in test_infer
generate_texts, token_nums, _ = pa_runner.infer(**infer_params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 390, in infer
generate_req(req_list_dummy, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager,
File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 1218, in generate_req
generate_token_with_clocking(model, cache_manager, batch, eplb_forwarder)
File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 892, in generate_token_with_clocking
res = generate_token(model, cache_manager, input_batch_in, eplb_forwarder)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 616, in generate_token
logits = model.forward(
^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/atb_llm/runner/model_runner.py", line 324, in forward
res = self.model.forward(**kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/atb_llm/models/qwen2_moe/flash_causal_qwen2_moe.py", line 438, in forward
logits = super().forward(**kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/Ascend/atb-models/atb_llm/models/base/flash_causal_lm.py", line 519, in forward
self.init_ascend_weight()
File "/usr/local/Ascend/atb-models/atb_llm/models/qwen2_moe/flash_causal_qwen2_moe.py", line 311, in init_ascend_weight
self.acl_encoder_operation.set_param(self.acl_param_encoder)
RuntimeError: Build model graph failed. enable log: export ASCEND_SLOG_PRINT_TO_STDOUT=1, export ASCEND_MODULE_LOG_LEVEL=ATB=3 to find the first error. For more details, see the MindIE official document.
总结:rank_table设置且不用dist.broadcast_object_list的时候一切正常,只要使用了broadcast_object_list,要么forward报异常,要么直接给卡死。不使用ranktable的时候也是直接报错
可能的原因:
1、双机 Asend HDK驱动版本不一致,主机 是 24.1.rc3; 从机 是 25.3.rc1;查看mindie官方的文档(2.2.RCX文档显示这两个驱动都是支持的)
2、torch_npu 的生态问题。torch 集合通信接口与底层的 HCCL 接口没有适配好。从而导致 torch 的 broadcast 与算子通信互相冲突;
容器中~下的所有日志文件我已经打包好,参考下面的附件。如果可以的话希望能够帮忙分析一下。
master_log.tar.gz
slave_log.tar.gz


@wangchao430 您好,我验证了下,我上次发您的程序确实是正常的,后来我尝试在mindie中的atb-models中使用dist.broadcast同步数据,结果发生报错/卡死。下面是复现过程,如果可以的话希望您能帮忙分析下!下面附上了详细的过程、时间、运行结果等!
准备文件:run_pa.py
相比源文件,添加了dist.init_process_group过程与test_infer、test_broadcast_data两个函数以及在main中调用。# Copyright Huawei Technologies Co., Ltd. 2023-2024. All rights reserved. import copy import json import math import os import time import torch.distributed as dist import torch import torch_npu from atb_llm.runner.model_runner import ModelRunner from atb_llm.utils.cpu_binding import NpuHbmInfo from atb_llm.utils.env import ENV from atb_llm.utils.log import logger, print_log from atb_llm.utils import file_utils from atb_llm.models.deepseekv2.eplb.eplb_planner.eplb_worker import EplbWorker from atb_llm.utils.moe_utils import EPLBType from atb_llm.utils.argument_utils import ArgumentParser, StringArgumentValidator, \ NumberArgumentValidator, MAX_JSON_LENGTH from examples.server.cache import CacheConfig, ModelConfig, CacheManager, HiddenStateCache from examples.server.generate import decode_token, generate_req from examples.server.request import request_from_token OPERATOR_BOUND = 576 MAX_KW_ARGS_LENGTH = 4096 PERF_FILE = "./examples/models/performance.csv" PERF_COLUMNS = "batch, input_len, output_len, embedding_len, first_token_time(ms), \ non_first_token_time(ms), ResponseTime(ms),E2E Throughput Average(Tokens/s)\n" PRED_FILE = "./examples/models/result.json" MICROBATCH_SIZE = "microbatch_size" class PARunner: def __init__(self, **kwargs): from datetime import timedelta dist.init_process_group( backend="hccl", world_size=kwargs.get('world_size', '1'), rank=kwargs.get('rank', '0'), timeout=timedelta(seconds=1800) ) logger.info(f"dgi: kwargs:{kwargs}") self.rank = kwargs.get('rank', '0') logger.info(f"Initialized HCCL process group for rank {self.rank}") self.local_rank = kwargs.get('local_rank', self.rank) self.world_size = kwargs.get('world_size', '1') self.model_path = kwargs.get('model_path', None) self.input_text = kwargs.get('input_text', None) self.max_prefill_batch_size = kwargs.get('max_prefill_batch_size', None) self.max_batch_size = kwargs.get('max_batch_size', None) self.max_input_length = kwargs.get('max_input_length', None) self.max_output_length = kwargs.get('max_output_length', None) self.max_position_embeddings = kwargs.get('max_position_embeddings', None) self.max_prefill_tokens = kwargs.get('max_prefill_tokens', None) self.enable_atb_torch = kwargs.get('enable_atb_torch', False) self.trust_remote_code = kwargs.get('trust_remote_code', False) self.block_size = kwargs.get('block_size', None) self.chat_template = kwargs.get('chat_template', None) self.load_tokenizer = kwargs.get('load_tokenizer', True) self.prefill_length = kwargs.get('prefill_length', 8192) self.distributed_enable = kwargs.get('distributed_enable', False) self.llm_config_path = kwargs.get('llm_config_path', None) self.check_limits() kw_args = kwargs.get('kw_args', '') if len(kw_args) > MAX_KW_ARGS_LENGTH: raise ValueError(f'The length of kw_args is too long and should be within (0, {MAX_KW_ARGS_LENGTH}]') if kw_args: try: kw_args = json.loads(kw_args) except json.decoder.JSONDecodeError: print_log(self.local_rank, logger.warning, 'Please checkout input of kw_args!') kw_args = {} else: kw_args = {} if kwargs.get(MICROBATCH_SIZE, -1) == -1: microbatch_size = self.max_batch_size elif kwargs.get(MICROBATCH_SIZE, -1) <= self.max_batch_size and kwargs.get(MICROBATCH_SIZE, -1) > 0: microbatch_size = kwargs.get(MICROBATCH_SIZE, -1) else: raise ValueError(f'Micro batch size should be within [1, {self.max_batch_size}] or -1!') if 'num_speculative_tokens' not in kw_args: kw_args['num_speculative_tokens'] = int(ENV.deepseek_mtp) self.model = ModelRunner( self.model_path, rank=self.rank, world_size=self.world_size, local_rank=self.local_rank, max_position_embeddings=self.max_position_embeddings, enable_atb_torch=self.enable_atb_torch, load_tokenizer=self.load_tokenizer, lora_modules=kwargs.get('lora_modules', None), max_loras=kwargs.get('max_loras', 0), max_lora_rank=kwargs.get('max_lora_rank', 0), trust_remote_code=self.trust_remote_code, llm_config_path=self.llm_config_path, dp=kwargs.get("dp", -1), tp=kwargs.get("tp", -1), moe_tp=kwargs.get("moe_tp", -1), pp=kwargs.get("pp", -1), microbatch_size=microbatch_size, moe_ep=kwargs.get("moe_ep", -1), sp=kwargs.get("sp", -1), cp=kwargs.get("cp", -1), distributed_enable=self.distributed_enable, max_batch_size=self.max_batch_size, **kw_args ) self.tokenizer = self.model.tokenizer if self.chat_template: self.tokenizer.chat_template = self._load_chat_template(self.chat_template) self.dtype = self.model.dtype self.quantize = self.model.quantize self.kv_quant_type = self.model.kv_quant_type self.fa_quant_type = self.model.fa_quant_type self.model.load_weights(**kw_args) self.device = self.model.device self.model_config = ModelConfig(self.model.num_heads, self.model.num_kv_heads, self.model.config.num_key_value_heads \ if hasattr(self.model.config, 'num_key_value_heads') \ else self.model.num_kv_heads, self.model.k_head_size, self.model.v_head_size, self.model.num_layers, self.model.device, self.model.dtype, self.model.soc_info, self.kv_quant_type, self.fa_quant_type, self.model.mapping, self.model.kvcache_quant_layers, getattr(self.model.config, 'cla_share_factor', 1), getattr(self.model.config, 'model_type', None), self.model.enable_nz) self.max_memory = NpuHbmInfo.get_hbm_capacity(self.local_rank, self.world_size, self.model.soc_info.need_nz) self.init_memory = int( self.max_memory * NpuHbmInfo.get_hbm_usage(self.local_rank, self.world_size, self.model.soc_info.need_nz)) print_log(self.rank, logger.info, f'hbm_capacity(GB): {self.max_memory / (1024 ** 3)}, ' f'init_memory(GB): {self.init_memory / (1024 ** 3)}') self.warm_up_memory = 0 self.warm_up_num_blocks = 0 self.cache_manager = None self.hidden_state_cache = None self.compress_head_enable = ENV.compress_head_enable self.omni_attention_enable = ENV.omni_attention_enable self.eplb_forwarder = None eplb_level = getattr(self.model.model, "eplb_level", 0) if eplb_level == EPLBType.DYNAMIC_EPLB: self.eplb_worker = EplbWorker(self.model, self.rank, self.model_path, self.device) self.eplb_forwarder = self.eplb_worker.eplb_forwarder def test_broadcast_data(self): if self.rank != 0: data = [None] else: data = [1] logger.info("before broadcast") dist.broadcast_object_list(data, src=0) logger.info("after broadcast") def __repr__(self): return ( "PARunner(" + f"model_path={self.model_path}, " + f"input_text={self.input_text}, " + f"max_position_embeddings={self.max_position_embeddings}, " + f"max_input_length={self.max_input_length}, " + f"max_output_length={self.max_output_length}, " + f"max_prefill_tokens={self.max_prefill_tokens}, " + f"load_tokenizer={self.load_tokenizer}, " + f"enable_atb_torch={self.enable_atb_torch}, " + f"max_prefill_batch_size={self.max_prefill_batch_size}, " + f"max_batch_size={self.max_batch_size}, " + f"dtype={self.dtype}, " + f"block_size={self.block_size}, " + f"model_config={self.model_config}, " + f"max_memory={self.max_memory}, " ) @staticmethod def save_input_output_ids(req_list): for idx, req in enumerate(req_list): input_ids_save_filename = f"input_ids_{idx}.pth" input_ids_save_path = os.path.join(ENV.token_ids_save_folder, input_ids_save_filename) input_ids_save_path = file_utils.standardize_path(input_ids_save_path) file_utils.check_file_safety(input_ids_save_path, 'w', is_check_file_size=False) torch.save(req.input_ids.cpu(), input_ids_save_path) os.chmod(input_ids_save_path, 0o640) output_ids_save_filename = f"output_ids_{idx}.txt" output_ids_save_path = os.path.join(ENV.token_ids_save_folder, output_ids_save_filename) with file_utils.safe_open(output_ids_save_path, 'w', encoding='utf-8') as f: f.write(' '.join(map(str, req.out_token_list))) @staticmethod def _update_inputs_items(inputs): input_texts, input_ids, input_conversations, input_adapters = [], [], [], [] if isinstance(inputs, list) and inputs: if isinstance(inputs[0], str): input_texts = inputs elif isinstance(inputs[0], torch.Tensor): input_ids = inputs elif isinstance(inputs[0], list) and inputs[0]: if isinstance(inputs[0][0], int): input_ids = inputs elif isinstance(inputs[0][0], dict): input_conversations = inputs elif isinstance(inputs[0], dict): prompt_key = "prompt" if isinstance(inputs[0].get(prompt_key), torch.Tensor): input_ids = [input.get(prompt_key) for input in inputs] else: input_texts = [input.get(prompt_key) for input in inputs] input_adapters = [input.get("adapter") for input in inputs] return input_texts, input_ids, input_conversations, input_adapters @staticmethod def _load_chat_template(chat_template: str): if os.path.exists(chat_template): with file_utils.safe_open(chat_template, 'r', encoding='utf-8') as f: chat_template_content = f.read() else: chat_template_content = chat_template return chat_template_content def check_limits(self): if self.max_input_length <= 0: print_log(self.rank, logger.warning, f"'max_input_length' must be greater than 0, but is {self.max_input_length}. " "Please increase the value of 'max_input_length'. ") if self.max_output_length <= 0: print_log(self.rank, logger.warning, f"'max_output_length' must be greater than 0, but is {self.max_output_length}. " "Please increase the value of 'max_output_length'. ") if self.max_position_embeddings and \ self.max_position_embeddings < self.max_input_length + self.max_output_length: print_log(self.rank, logger.warning, f"'max_position_embeddings' must be no less than " f"{self.max_input_length + self.max_output_length}, but is {self.max_position_embeddings}. " "Please increase the value of 'max_position_embeddings'. ") if self.max_prefill_batch_size and self.max_prefill_batch_size <= 0: print_log(self.rank, logger.warning, f"'max_prefill_batch_size' must be greater than 0, but is {self.max_prefill_batch_size}. " "Please increase the value of 'max_prefill_batch_size'. ") if self.max_batch_size <= 0: print_log(self.rank, logger.warning, f"'max_batch_size' must be greater than 0, but is {self.max_batch_size}. " "Please increase the value of 'max_batch_size'. ") if self.block_size > OPERATOR_BOUND: print_log(self.rank, logger.warning, f"'block_size' must be lower than or equal to 576, but is {self.block_size}. " "Please decrease the value of 'block_size'. ") if self.block_size <= 0: print_log(self.rank, logger.warning, f"'block_size' is too small: {self.block_size}. " "Please increase the value of 'block_size'. ") if self.max_prefill_tokens != -1 and \ self.max_prefill_tokens < self.max_input_length: print_log(self.rank, logger.warning, f"'max_prefill_tokens' must be greater than or equal to {self.max_input_length}, " f"but is {self.max_prefill_tokens}. Please increse the value of 'max_prefill_tokens'. ") def calc_max_length(self): if self.model.mapping.has_attn_cp() and self.model.mapping.has_attn_inner_sp(): max_length = math.ceil((math.ceil(self.max_input_length / self.model.mapping.attn_cp.group_size) + self.max_output_length) / self.model.mapping.attn_inner_sp.group_size) elif self.model.mapping.has_attn_inner_sp(): max_length = math.ceil((self.max_input_length + self.max_output_length) / self.model.mapping.attn_inner_sp.group_size) elif self.model.mapping.has_attn_cp(): # CP will perform padding on the input. max_length = (math.ceil(self.max_input_length / self.model.mapping.attn_cp.group_size) + math.ceil(self.max_output_length / self.model.mapping.attn_cp.group_size)) else: max_length = self.max_input_length + self.max_output_length return max_length def warm_up(self): self.model.model.warmup_is_end = False if self.max_prefill_batch_size is None: self.max_prefill_batch_size = self.max_batch_size if ENV.modeltest_pd_split_enable: if self.max_input_length > self.prefill_length: self.max_prefill_batch_size = 1 else: self.max_prefill_batch_size = math.floor(self.prefill_length / self.max_input_length) if self.max_prefill_batch_size > self.max_batch_size: self.max_prefill_batch_size = self.max_batch_size effective_max_prefill_batch_size = math.ceil( self.max_prefill_batch_size / self.model.mapping.attn_dp.group_size) effective_max_batch_size = math.ceil( self.max_batch_size / self.model.mapping.attn_dp.group_size) if self.max_prefill_tokens == -1: self.max_prefill_tokens = effective_max_prefill_batch_size * (self.max_input_length + self.max_output_length) print_log(self.rank, logger.info, "---------------begin warm_up---------------") try: max_length = self.calc_max_length() self.warm_up_num_blocks = math.ceil(max_length / self.block_size) * effective_max_batch_size except ZeroDivisionError as e: raise ZeroDivisionError from e cache_config = CacheConfig(self.warm_up_num_blocks, self.block_size) if self.compress_head_enable or self.omni_attention_enable: cache_config = CacheConfig(self.warm_up_num_blocks, self.block_size, \ self.max_input_length, self.max_output_length, effective_max_batch_size, self.rank, self.world_size) self.cache_manager = CacheManager(cache_config, self.model_config) if ENV.deepseek_mtp: batch_size_for_cache = max(self.max_batch_size, self.max_prefill_batch_size) self.cache_manager.hidden_state_cache = HiddenStateCache(self.model.config.hidden_size, ENV.deepseek_mtp, batch_size_for_cache, self.model.dtype, self.cache_manager.device) else: self.hidden_state_cache = None self.model.postprocessor.max_new_tokens = 1 input_ids = torch.ones(self.max_input_length, dtype=torch.int64) req_list = [ request_from_token(input_ids, 1, self.block_size, req_idx=idx, mapping=self.model.mapping) for idx in range(self.max_prefill_batch_size) ] generate_req(req_list, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager, self.eplb_forwarder) self.warm_up_memory = int( self.max_memory * NpuHbmInfo.get_hbm_usage(self.local_rank, self.world_size, self.model.soc_info.need_nz)) print_log(self.rank, logger.info, f'warmup_memory(GB): {self.warm_up_memory / (1024 ** 3): .2f}') print_log(self.rank, logger.info, "---------------end warm_up---------------") self.model.model.warmup_is_end = True def infer(self, inputs, batch_size, max_output_length, ignore_eos, is_chat_model=False, **kwargs): if torch.distributed.is_initialized(): logger.info(">>>> torch.distributed is already initialized.") else: logger.error(">>>> torch.distributed has not been initialized!") print_log(self.rank, logger.info, "---------------begin inference---------------") if ignore_eos: self.model.postprocessor.eos_token_id = [] is_truncation = kwargs.get("truncation", False) input_ids, adapter_ids = self._build_model_inputs(inputs, is_chat_model, is_truncation) if len(input_ids) == 1: req_list = [request_from_token(input_ids[0], max_output_length, self.block_size, req_idx=idx, adapter_id=adapter_ids[0], mapping=self.model.mapping, postprocessor=self.model.postprocessor) for idx in range(batch_size)] else: req_list = [request_from_token(input_ids_ins, max_output_length, self.block_size, req_idx=idx, adapter_id=adapter_id, mapping=self.model.mapping, postprocessor=self.model.postprocessor) for idx, (input_ids_ins, adapter_id) in enumerate(zip(input_ids, adapter_ids))] print_log(self.rank, logger.debug, f'req_list[0].input_ids: {req_list[0].input_ids}') if not self.cache_manager: if self.max_prefill_tokens == -1: self.max_prefill_tokens = self.max_batch_size * (self.max_input_length + self.max_output_length) cache_block_size = self.block_size * self.model.num_kv_heads * self.model.head_size dtype_size = CacheManager.get_dtype_size(self.dtype) total_cache_size = self.model.num_layers * cache_block_size * 2 * dtype_size max_memory = ENV.memory_fraction * self.max_memory free_memory = max_memory - ENV.reserved_memory_gb * (1 << 30) - ( self.warm_up_memory if self.warm_up_memory != 0 else self.init_memory) print_log(self.rank, logger.info, f"infer max_memory(GB): {max_memory / (1024 ** 3): .2f}, " f"warm_up_memory(GB): {self.warm_up_memory / (1024 ** 3): .2f}, " f"free_memory(GB): {free_memory / (1024 ** 3): .2f}") try: num_blocks = int(free_memory // total_cache_size) except ZeroDivisionError as e: raise ZeroDivisionError("total_cache_size should not be 0") from e print_log(self.rank, logger.info, f"num_blocks: {num_blocks}, free_memory: {free_memory}") cache_config = CacheConfig(num_blocks, self.block_size) if self.compress_head_enable: cache_config = CacheConfig(self.warm_up_num_blocks, self.block_size, \ self.max_input_length, self.max_output_length, self.max_batch_size, self.rank, self.world_size) self.cache_manager = CacheManager(cache_config, self.model_config) if ENV.benchmark_enable: req_list_dummy = copy.deepcopy(req_list) self.model.postprocessor.max_new_tokens = 2 generate_req(req_list_dummy, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager, self.eplb_forwarder) self.model.postprocessor.max_new_tokens = max_output_length skip_special_tokens = kwargs.get("skip_special_tokens", False) if not ENV.profiling_enable: print_log(self.rank, logger.debug, "No profiling") torch.npu.synchronize() e2e_start = time.time() generate_req(req_list, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager, self.eplb_forwarder) _, _ = decode_token(req_list, self.tokenizer, skip_special_tokens) torch.npu.synchronize() e2e_end = time.time() e2e_time = e2e_end - e2e_start else: print_log(self.rank, logger.debug, "Enter profiling") profiling_path = ENV.profiling_filepath profiler_level = torch_npu.profiler.ProfilerLevel actual_profiler_level = getattr(profiler_level, ENV.profiling_level) torch.npu.synchronize() e2e_start = time.time() experimental_config = torch_npu.profiler._ExperimentalConfig( aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, profiler_level=actual_profiler_level, l2_cache=False, data_simplification=False ) with torch_npu.profiler.profile( activities=[ torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU ], on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profiling_path), record_shapes=True, profile_memory=True, with_stack=False, with_flops=False, with_modules=False, experimental_config=experimental_config): generate_req(req_list, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager, self.eplb_forwarder) torch.npu.synchronize() e2e_end = time.time() e2e_time = e2e_end - e2e_start generate_text_list, token_num_list = decode_token(req_list, self.tokenizer, skip_special_tokens) if ENV.modeltest_dataset_specified: ENV.update() if ENV.token_ids_save_enable: if self.model.mapping.has_pp(): if self.rank == self.world_size - 1: self.save_input_output_ids(req_list) else: if self.local_rank == 0: self.save_input_output_ids(req_list) print_log(self.rank, logger.info, "---------------end inference---------------") return generate_text_list, token_num_list, e2e_time def shutdown(self): eplb_level = getattr(self.model.config, "eplb_level", 0) if eplb_level == EPLBType.DYNAMIC_EPLB: self.eplb_worker.eplb_loader.shutdown() self.eplb_worker.eplb_planner.shutdown() def _build_model_inputs(self, inputs, is_chat_model, is_truncation=False): input_texts, input_ids, input_conversations, input_adapters = self._update_inputs_items(inputs) if not (input_texts or input_ids or input_conversations): raise ValueError(f"The inputs of `PARunner.infer` must be as List[str], List[torch.Tensor], List[List[int]]" f" or List[List[Dict]]. Now the inputs ({inputs}) is not acceptable or is empty.") if not input_ids and not self.tokenizer: raise ValueError("Can‘t encode input_texts, because load_tokenizer=False, " "please enter input_ids or enable load_tokenizer") if is_chat_model: if input_conversations: input_ids = self.model.build_inputs(input_conversations) elif input_texts: input_conversations = [[{"role": "user", "content": t}] for t in input_texts] input_ids = self.model.build_inputs(input_conversations) else: print_log(self.rank, logger.warning, "Neither conversations nor input_texts exist, " "'chat' parameter is not effective.") elif input_texts: input_ids = [self.tokenizer([text], return_tensors="pt", truncation=is_truncation)["input_ids"].flatten() for text in input_texts] # token长度校验 for item in input_ids: if len(item) > self.max_input_length: print_log(self.rank, logger.warning, "Num of tokens in input_text is larger than max_input_length. " "Please shorten input_text to avoid out of memory.") if not input_adapters: input_adapters = [None] * len(input_ids) return input_ids, input_adapters def cmd_bool(cmd_arg): if cmd_arg == "True": return True elif cmd_arg == "False": return False raise ValueError(f"{cmd_arg} should be a boolean") def parse_list_of_json(list_json): return json.loads(list_json) def parse_ids(list_str): return [int(item) for item in list_str.split(',')] def input_texts_parser(value): if os.path.isfile(value): with file_utils.safe_open(value, 'r') as opened_file: return opened_file.read() else: return value #define Argument Parser def parse_arguments(): store_true = 'store_true' parser = ArgumentParser() parser.add_argument('--model_path', type=str, help="model and tokenizer path") parser.add_argument( '--input_texts', type=input_texts_parser, nargs='+', default=["What's deep learning?"]) parser.add_argument( '--input_ids', type=parse_ids, nargs='+', default=None) parser.add_argument( '--input_file', type=str, help='This parameter is used to input multi-turn dialogue information in the form ' 'of a jsonl file, with each line in the format of a List[Dict]. Each dictionary ' '(Dict) must contain at least two fields: "role" and "content".', default=None) parser.add_argument( '--input_dict', help="Lora input, accepted format: " "'[{\"prompt\": \"prompt in text\", \"adapter\": \"adapter id defined in lora_adapter param\"}]'", type=parse_list_of_json, default=None) parser.add_argument("--max_prefill_batch_size", type=int, default=None) parser.add_argument("--max_batch_size", type=int, default=1) parser.add_argument('--max_input_length', type=int, default=1024) parser.add_argument('--max_output_length', type=int, default=20) parser.add_argument('--max_position_embeddings', type=int, default=None) parser.add_argument('--max_prefill_tokens', type=int, default=-1) parser.add_argument("--block_size", type=int, default=128) parser.add_argument('--chat_template', type=str, default=None) parser.add_argument('--ignore_eos', action=store_true) parser.add_argument('--is_chat_model', action=store_true) parser.add_argument('--is_embedding_model', action=store_true) parser.add_argument('--load_tokenizer', type=cmd_bool, default=True) parser.add_argument('--enable_atb_torch', action=store_true) parser.add_argument('--trust_remote_code', action=store_true) parser.add_argument('--dp', type=int, default=-1) parser.add_argument('--tp', type=int, default=-1) parser.add_argument('--moe_tp', type=int, default=-1) parser.add_argument('--pp', type=int, default=-1) parser.add_argument('--microbatch_size', type=int, default=-1) parser.add_argument('--moe_ep', type=int, default=-1) parser.add_argument('--sp', type=int, default=-1) parser.add_argument('--cp', type=int, default=-1) parser.add_argument('--kw_args', type=str, default='', help='json input') parser.add_argument('--lora_modules', type=str, default=None, help="Define lora modules, accepted format:" "'{\"adapter1\": \"/path/to/lora1\", \"adapter2\": \"/path/to/lora2\"}'", validator=StringArgumentValidator(min_length=2, max_length=MAX_JSON_LENGTH, allow_none=True)) parser.add_argument('--max_loras', type=int, default=0, help="Define the MAX number of loRAs to store in NPU memory", validator=NumberArgumentValidator(min_value=0)) parser.add_argument('--max_lora_rank', type=int, default=0, help="Define the MAX loRA rank", validator=NumberArgumentValidator(min_value=0)) parser.add_argument('--distributed_enable', action=store_true) return parser.parse_args() def test_infer(pa_runner): infer_params = { "inputs": ["What's deep learning?"], "batch_size": 1, "max_output_length": 10, "ignore_eos": True, "is_chat_model": True } generate_texts, token_nums, _ = pa_runner.infer(**infer_params) length = len(infer_inputs) for i, generate_text in enumerate(generate_texts): if input_dict['is_embedding_model']: embedding_tensor_path = f"{os.getcwd()}/examples/embedding_tensor" print_log(rank, logger.info, f"Context[{i}]: \nembedding tensor path is: {embedding_tensor_path}") else: print_log(rank, logger.info, f'Answer[{i}]: {generate_text}') print_log(rank, logger.info, f'Generate[{i}] token num: {token_nums[i]}') if __name__ == '__main__': args = parse_arguments() rank = ENV.rank local_rank = ENV.local_rank world_size = ENV.world_size input_dict = { 'rank': rank, 'world_size': world_size, 'local_rank': local_rank, **vars(args) } # 输入优先级input_dict > input_ids > input_texts if args.input_dict: infer_inputs = args.input_dict elif args.input_ids: infer_inputs = args.input_ids else: infer_inputs = args.input_texts if args.is_chat_model and args.input_file: conversations = [] with file_utils.safe_open(args.input_file, 'r', encoding='utf-8') as file: for line in file_utils.safe_readlines(file): data_line = json.loads(line) conversations.append(data_line) infer_inputs = conversations pa_runner = PARunner(**input_dict) print_log(rank, logger.info, f'pa_runner: {pa_runner}') # pa_runner.test_broadcast_data() print_log(rank, logger.info, f'before_test_infer') test_infer(pa_runner) test_infer(pa_runner) print_log(rank, logger.info, f'after_test_infer') print_log(rank, logger.info, f'before_test_infer1') test_infer(pa_runner) print_log(rank, logger.info, f'after_test_infer1') # exit(0) pa_runner.warm_up() infer_params = { "inputs": infer_inputs, "batch_size": args.max_batch_size, "max_output_length": args.max_output_length, "ignore_eos": args.ignore_eos, "is_chat_model": args.is_chat_model } generate_texts, token_nums, _ = pa_runner.infer(**infer_params) length = len(infer_inputs) for i, generate_text in enumerate(generate_texts): if input_dict['is_embedding_model']: embedding_tensor_path = f"{os.getcwd()}/examples/embedding_tensor" print_log(rank, logger.info, f"Context[{i}]: \nembedding tensor path is: {embedding_tensor_path}") else: print_log(rank, logger.info, f'Answer[{i}]: {generate_text}') print_log(rank, logger.info, f'Generate[{i}] token num: {token_nums[i]}')# 下面的命令两台主机都要执行,注意:1、run_pa的挂载(容器内修改文件也可)2.rank_table_file的挂载 docker run -it \ --name dgi-dev-mindie \ --network=host \ --privileged=true \ --device=/dev/davinci_manager \ --device=/dev/hisi_hdc \ --device=/dev/devmm_svm \ -v /usr/local/lib/node_modules/:/usr/local/lib/node_modules/:ro \ -v /usr/local/bin/node:/usr/local/bin/node:ro \ -v /usr/local/dcmi:/usr/local/dcmi \ -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ -v /usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64/common \ -v /usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/driver \ -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ -v /etc/ascend_install.info:/etc/ascend_install.info \ -v /etc/vnpu.cfg:/etc/vnpu.cfg \ -v /usr/local/sbin:/usr/local/sbin \ -v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \ -v /data1/:/data1/ \ -v /data2/:/data2/ \ -v /data1/MindIE_Auto_Deploy_Tool/rank_table_file.json:/data1/MindIE_Auto_Deploy_Tool/rank_table_file.json \ -v /data1/atb-models/examples/run_pa.py:/usr/local/Ascend/atb-models/examples/run_pa.py \ -w /root \ -u root \ --shm-size=64GB \ --ulimit msgqueue=-1 \ -e ASCEND_RT_VISIBLE_DEVICES="0,1,2,3,4,5,6,7" \ swr.cn-south-1.myhuaweicloud.com/ascendhub/mindie:2.2.RC1-800I-A2-py311-openeuler24.03-lts # 进容器后执行下面的命令 source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/atb-models/set_env.sh export ATB_LLM_HCCL_ENABLE=1 export ATB_LLM_COMM_BACKEND="hccl" export HCCL_CONNECT_TIMEOUT=7200 export HCCL_EXEC_TIMEOUT=0 export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 export ATB_LLM_BENCHMARK_ENABLE=1 export ATB_LLM_ENABLE_AUTO_TRANSPOSE=0 # 打开 ATB 日志 # export ASCEND_GLOBAL_LOG_LEVEL=3 # export ASCEND_SLOG_PRINT_TO_STDOUT=1 # MINDIE日志 export MINDIE_LOG_TO_STDOUT=1 export HCCL_DETERMINISTIC=false export LCCL_DETERMINISTIC=0 export HCCL_BUFFSIZE=120 export ATB_WORKSPACE_MEM_ALLOC_GLOBAL=1 # 使用ranktable! export RANK_TABLE_FILE=/data1/MindIE_Auto_Deploy_Tool/rank_table_file.json首先是使用rank_table推理,不用broadcast
主机执行:
# 主机执行 torchrun \ --nproc_per_node 8 \ --master_port 12345 \ --nnodes 2 \ --node_rank 0 \ --master_addr 10.17.44.64 \ /usr/local/Ascend/atb-models/examples/run_pa.py \ --model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/master.log从机执行:
# 从机执行 torchrun \ --nproc_per_node 8 \ --master_port 12345 \ --nnodes 2 \ --node_rank 1 \ --master_addr 10.17.44.64 \ /usr/local/Ascend/atb-models/examples/run_pa.py \ --model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/slave.log2026-01-25 19:55启动
日志中可能报警告:[2026-01-25 19:55:53.460882] [833] [833] [llmmodels] [WARN] [context_factory.cpp:37] Failed to aclInit: 100002但结果正常推理,输出
[2026-01-25 19:56:42,324] [829] [281473611860544] [llmmodels] [INFO] [run_pa.py-652] : Answer[0]: What's the difference between deep learning and machine learning? What's the difference between deep learning and neural [2026-01-25 19:56:42,324] [829] [281473611860544] [llmmodels] [INFO] [run_pa.py-653] : Generate[0] token num: (0, 20)使用rank_table推理,推理前加入broadcast
两台容器同时进
vim /usr/local/Ascend/atb-models/examples/run_pa.py +626
取消注释pa_runner.test_broadcast_data()
主机执行:# 主机执行 torchrun \ --nproc_per_node 8 \ --master_port 12345 \ --nnodes 2 \ --node_rank 0 \ --master_addr 10.17.44.64 \ /usr/local/Ascend/atb-models/examples/run_pa.py \ --model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/master1.log从机执行:
# 从机执行 torchrun \ --nproc_per_node 8 \ --master_port 12345 \ --nnodes 2 \ --node_rank 1 \ --master_addr 10.17.44.64 \ /usr/local/Ascend/atb-models/examples/run_pa.py \ --model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/slave1.log2026-01-25 20:01启动,报错:
[2026-01-25 20:02:05,455] [3754] [281473078336064] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory /usr/local/lib64/python3.11/site-packages/torch_npu/utils/storage.py:41: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly. To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage() if self.device.type != 'cpu': [2026-01-25 20:02:05,521] [3754] [281473078336064] [llmmodels] [INFO] [generate.py-1130] : ------total req num: 1, infer start-------- [2026-01-25 20:02:05,521] [3758] [281473476311616] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:02:06,782] [3755] [281473268972096] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:02:07,495] [3760] [281473129114176] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:02:08,468] [3761] [281472942824000] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:02:08,474] [3757] [281473277147712] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:02:10,166] [3759] [281473160243776] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:02:13,227] [3756] [281473215760960] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:02:15,449] [3759] [281473160243776] [llmmodels] [INFO] [run_pa.py-165] : after broadcast [2026-01-25 20:02:15,449] [3755] [281473268972096] [llmmodels] [INFO] [run_pa.py-165] : after broadcast [2026-01-25 20:02:15,449] [3758] [281473476311616] [llmmodels] [INFO] [run_pa.py-165] : after broadcast [2026-01-25 20:02:15,450] [3759] [281473160243776] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized. [2026-01-25 20:02:15,450] [3760] [281473129114176] [llmmodels] [INFO] [run_pa.py-165] : after broadcast [2026-01-25 20:02:15,450] [3755] [281473268972096] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized. [2026-01-25 20:02:15,450] [3758] [281473476311616] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized. [2026-01-25 20:02:15,450] [3760] [281473129114176] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized. [2026-01-25 20:02:15,450] [3757] [281473277147712] [llmmodels] [INFO] [run_pa.py-165] : after broadcast [2026-01-25 20:02:15,450] [3757] [281473277147712] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized. [2026-01-25 20:02:15,454] [3756] [281473215760960] [llmmodels] [INFO] [run_pa.py-165] : after broadcast [2026-01-25 20:02:15,454] [3756] [281473215760960] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized. [2026-01-25 20:02:15,455] [3761] [281472942824000] [llmmodels] [INFO] [run_pa.py-165] : after broadcast [2026-01-25 20:02:15,455] [3761] [281472942824000] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized. [2026-01-25 20:02:15,553] [3759] [281473160243776] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory [2026-01-25 20:02:15,555] [3755] [281473268972096] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory [2026-01-25 20:02:15,556] [3761] [281472942824000] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory [2026-01-25 20:02:15,562] [3758] [281473476311616] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory [2026-01-25 20:02:15,564] [3756] [281473215760960] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory [2026-01-25 20:02:15,565] [3757] [281473277147712] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory [2026-01-25 20:02:15,567] [3760] [281473129114176] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory mki_log mkdir /root/ascend/log/atb Traceback (most recent call last): File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 628, in <module> test_infer(pa_runner) File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 584, in test_infer generate_texts, token_nums, _ = pa_runner.infer(**infer_params) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 390, in infer generate_req(req_list_dummy, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager, File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 1218, in generate_req generate_token_with_clocking(model, cache_manager, batch, eplb_forwarder) File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 892, in generate_token_with_clocking res = generate_token(model, cache_manager, input_batch_in, eplb_forwarder) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 616, in generate_token logits = model.forward( ^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/atb_llm/runner/model_runner.py", line 324, in forward res = self.model.forward(**kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/atb_llm/models/qwen2_moe/flash_causal_qwen2_moe.py", line 438, in forward logits = super().forward(**kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/atb_llm/models/base/flash_causal_lm.py", line 519, in forward self.init_ascend_weight() File "/usr/local/Ascend/atb-models/atb_llm/models/qwen2_moe/flash_causal_qwen2_moe.py", line 311, in init_ascend_weight self.acl_encoder_operation.set_param(self.acl_param_encoder) RuntimeError: External Comm Manager: Create the hccl communication group failed. export ASCEND_GLOBAL_LOG_LEVEL=3, export ASCEND_SLOG_PRINT_TO_STDOUT=1 to see more details. Default log path is $HOME/atb/log. Traceback (most recent call last): File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 628, in <module> test_infer(pa_runner) File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 584, in test_infer generate_texts, token_nums, _ = pa_runner.infer(**infer_params) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 390, in infer generate_req(req_list_dummy, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager, File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 1218, in generate_req generate_token_with_clocking(model, cache_manager, batch, eplb_forwarder) File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 892, in generate_token_with_clocking res = generate_token(model, cache_manager, input_batch_in, eplb_forwarder) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 616, in generate_token logits = model.forward( ^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/atb_llm/runner/model_runner.py", line 324, in forward res = self.model.forward(**kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/atb_llm/models/qwen2_moe/flash_causal_qwen2_moe.py", line 438, in forward logits = super().forward(**kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/atb_llm/models/base/flash_causal_lm.py", line 519, in forward self.init_ascend_weight() File "/usr/local/Ascend/atb-models/atb_llm/models/qwen2_moe/flash_causal_qwen2_moe.py", line 311, in init_ascend_weight self.acl_encoder_operation.set_param(self.acl_param_encoder) RuntimeError: External Comm Manager: Create the hccl communication group failed. export ASCEND_GLOBAL_LOG_LEVEL=3, export ASCEND_SLOG_PRINT_TO_STDOUT=1 to see more details. Default log path is $HOME/atb/log.使用rank_table推理,推理后尝试broadcast:
两台容器同时进
vim /usr/local/Ascend/atb-models/examples/run_pa.py +626
将626行的pa_runner.test_broadcast_data()移动628行到两个test_infer之间,像是这样:624 pa_runner = PARunner(**input_dict) 625 print_log(rank, logger.info, f'pa_runner: {pa_runner}') 626 print_log(rank, logger.info, f'before_test_infer') 627 test_infer(pa_runner) 628 pa_runner.test_broadcast_data() 629 test_infer(pa_runner) 630 print_log(rank, logger.info, f'after_test_infer')主机执行:
# 主机执行 torchrun \ --nproc_per_node 8 \ --master_port 12345 \ --nnodes 2 \ --node_rank 0 \ --master_addr 10.17.44.64 \ /usr/local/Ascend/atb-models/examples/run_pa.py \ --model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/master2.log从机执行:
# 从机执行 torchrun \ --nproc_per_node 8 \ --master_port 12345 \ --nnodes 2 \ --node_rank 1 \ --master_addr 10.17.44.64 \ /usr/local/Ascend/atb-models/examples/run_pa.py \ --model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/slave2.log2026-01-25 20:08 启动,第一个test_infer正常,程序卡死在before_broadcast与after_broadcast之间,日志例如(从机,主机同样):等了十分钟还是没响应,ctrl+c停掉了。
[2026-01-25 20:09:56,791] [5677] [281473705536064] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:09:56,792] [5676] [281472899271232] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:09:56,792] [5678] [281473714686528] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:09:56,792] [5672] [281473198320192] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:09:56,792] [5674] [281472934701632] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:09:56,792] [5675] [281473338534464] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:09:56,792] [5673] [281473130969664] [llmmodels] [INFO] [run_pa.py-163] : before broadcast [2026-01-25 20:09:56,792] [5679] [281473051490880] [llmmodels] [INFO] [run_pa.py-163] : before broadcast在上一个的基础上不使用rank_table推理
紧接上一个例子,两个test_infer之间使用broadcast,并且unset RANK_TABLE_FILE变量
unset RANK_TABLE_FILE主机执行:
# 主机执行 torchrun \ --nproc_per_node 8 \ --master_port 12345 \ --nnodes 2 \ --node_rank 0 \ --master_addr 10.17.44.64 \ /usr/local/Ascend/atb-models/examples/run_pa.py \ --model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/master3.log从机执行:
# 从机执行 torchrun \ --nproc_per_node 8 \ --master_port 12345 \ --nnodes 2 \ --node_rank 1 \ --master_addr 10.17.44.64 \ /usr/local/Ascend/atb-models/examples/run_pa.py \ --model_path /data1/models/Qwen3-30B-A3B-Instruct-2507 2>&1 | tee ~/slave3.log2026-01-25 20:20启动,报错变成了:
[2026-01-25 20:21:21,470] [7981] [281473698683456] [llmmodels] [INFO] [generate.py-1130] : ------total req num: 1, infer start-------- [2026-01-25 20:21:21,969] [7985] [281473459296832] [llmmodels] [INFO] [run_pa.py-342] : >>>> torch.distributed is already initialized. [2026-01-25 20:21:22,102] [7985] [281473459296832] [llmmodels] [INFO] [cache.py-154] : kv cache will allocate 53.3173828125GB memory [2026-01-25 20:21:22.214388] [7981] [7981] [llmmodels] [INFO] [moe_decoder_model.cpp:211] model enableSplitFuse: 0 Traceback (most recent call last): File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 627, in <module> test_infer(pa_runner) File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 584, in test_infer generate_texts, token_nums, _ = pa_runner.infer(**infer_params) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/examples/run_pa.py", line 390, in infer generate_req(req_list_dummy, self.model, self.max_batch_size, self.max_prefill_tokens, self.cache_manager, File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 1218, in generate_req generate_token_with_clocking(model, cache_manager, batch, eplb_forwarder) File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 892, in generate_token_with_clocking res = generate_token(model, cache_manager, input_batch_in, eplb_forwarder) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/examples/server/generate.py", line 616, in generate_token logits = model.forward( ^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/atb_llm/runner/model_runner.py", line 324, in forward res = self.model.forward(**kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/atb_llm/models/qwen2_moe/flash_causal_qwen2_moe.py", line 438, in forward logits = super().forward(**kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/Ascend/atb-models/atb_llm/models/base/flash_causal_lm.py", line 519, in forward self.init_ascend_weight() File "/usr/local/Ascend/atb-models/atb_llm/models/qwen2_moe/flash_causal_qwen2_moe.py", line 311, in init_ascend_weight self.acl_encoder_operation.set_param(self.acl_param_encoder) RuntimeError: Build model graph failed. enable log: export ASCEND_SLOG_PRINT_TO_STDOUT=1, export ASCEND_MODULE_LOG_LEVEL=ATB=3 to find the first error. For more details, see the MindIE official document.总结:rank_table设置且不用dist.broadcast_object_list的时候一切正常,只要使用了broadcast_object_list,要么forward报异常,要么直接给卡死。不使用ranktable的时候也是直接报错
可能的原因:
1、双机 Asend HDK驱动版本不一致,主机 是 24.1.rc3; 从机 是 25.3.rc1;查看mindie官方的文档(2.2.RCX文档显示这两个驱动都是支持的)
2、torch_npu 的生态问题。torch 集合通信接口与底层的 HCCL 接口没有适配好。从而导致 torch 的 broadcast 与算子通信互相冲突;
容器中~下的所有日志文件我已经打包好,参考下面的附件。如果可以的话希望能够帮忙分析一下。
master_log.tar.gz
slave_log.tar.gz
- 根据当前plog的报错,显示是有“hccl multiple initialization”,对应表示同一个hccl通信域初始化了2次。怀疑有残留进程没杀干净。
- 不同的HDK版本在底层通信能力的支持上存在区别,无法完全保证能兼容的。 推荐升级至同样的驱动版本进行测试验证,并保证残留进程杀干净。


在提交新问题之前,请确保您已经在社区中搜索过相关问题,并使用了社区中提供的资源/工具后,仍未找到满意的解决方式。
环境信息
硬件:16卡910B2
使用场景及问题
torch_npu是否支持多机多卡broadcast数据??我没有找到文档示例,而且我的示例代码也没有跑通。
建联方式:ranktable建链,已经验证npu是通的。已经验证硬件使用mindie镜像可以多机推理。下面的操作是在quay.io/ascend/cann:8.3.rc2-910b-ubuntu22.04-py3.11为基础镜像的镜像里操作的:
其中:venvs/xxx/bin/activate是使用uv手动创建的虚拟环境,里面装了python=3.11以及torch=2.1.0、torchnpu=2.1.0.post13以及必要的依赖库。
进入容器后操作:
env:
ranktable:
代码:
双机16卡启动脚本:
欢迎加入社区,感谢您对社区的贡献 🎉!