36e450eb0f
we recently added disaggregated_prefill and ascend_forward_context feature byba3dfbd59eanddf0ec55162. This PR fix some nit introduced by them to make the code clear. 1. drop `current_platform` usage. It'll lead unknown circular import error in some case 2. update `set_ascend_forward_context` function to make the logic clear. for example, remove V0 support in this function. 3. Remove useless `self.local_rank_across_dp` in worker 4. Remove `soc_info.py` to use `get_ascend_soc_version` instead. - vLLM version: v0.10.0 - vLLM main:02f82fe438Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
123 lines
4.3 KiB
Python
123 lines
4.3 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
|
|
import torch.distributed as dist
|
|
|
|
from vllm_ascend.utils import AscendSocVersion, init_ascend_soc_version, get_ascend_soc_version
|
|
|
|
parser = argparse.ArgumentParser(
|
|
description="Arguments of rank table generator", )
|
|
parser.add_argument("--local-host", type=str, required=True, help="local ip")
|
|
parser.add_argument("--prefill-device-cnt",
|
|
type=int,
|
|
required=True,
|
|
help="number of prefill devices")
|
|
parser.add_argument("--decode-device-cnt",
|
|
type=int,
|
|
required=True,
|
|
help="number of decode devices")
|
|
args = parser.parse_args()
|
|
local_host = args.local_host
|
|
prefill_device_cnt = args.prefill_device_cnt
|
|
decode_device_cnt = args.decode_device_cnt
|
|
|
|
print("enter py")
|
|
|
|
hccn_tool_path = os.environ.get("HCCN_TOOL_PATH",
|
|
"/usr/local/Ascend/driver/tools/hccn_tool")
|
|
master_addr = os.environ.get("MASTER_ADDR")
|
|
master_port = os.environ.get("MASTER_PORT")
|
|
rank = os.environ.get("RANK")
|
|
local_rank = os.environ.get("LOCAL_RANK")
|
|
# This variable is set by torchrun,
|
|
# and is different from WORLD_SIZE in gen_rank_table.sh.
|
|
world_size = os.environ.get("WORLD_SIZE")
|
|
|
|
init_ascend_soc_version()
|
|
soc_info = get_ascend_soc_version()
|
|
|
|
|
|
def get_cmd_stdout(cmd):
|
|
import subprocess
|
|
return subprocess.run(cmd, capture_output=True,
|
|
shell=True).stdout.decode("utf-8").strip()
|
|
|
|
|
|
print(f"local_host: {local_host}")
|
|
print("gen ranktable.json")
|
|
|
|
num_cards = get_cmd_stdout("npu-smi info -l | grep \"Total Count\"").split(
|
|
":")[1].strip()
|
|
num_cards = int(num_cards)
|
|
chips_per_card = get_cmd_stdout("npu-smi info -l | grep \"Chip Count\"").split(
|
|
"\n")[0].split(":")[1].strip()
|
|
chips_per_card = int(chips_per_card)
|
|
|
|
# generate local device list for local rank 0, and gather it to all ranks
|
|
local_device_list: list[dict[str, str]] = list()
|
|
if local_rank == "0":
|
|
super_pod_id = "0"
|
|
for card_id in range(num_cards):
|
|
for chip_id in range(chips_per_card):
|
|
device_id = card_id * chips_per_card + chip_id
|
|
if soc_info == AscendSocVersion.A3:
|
|
device_ip = get_cmd_stdout(
|
|
f"{hccn_tool_path} -i {device_id} -vnic -g | grep ipaddr"
|
|
).split(":")[1].strip()
|
|
super_device_id = get_cmd_stdout(
|
|
f"npu-smi info -t spod-info -i {card_id} -c {chip_id} | grep SDID"
|
|
).split(":")[1].strip()
|
|
super_pod_id = get_cmd_stdout(
|
|
f"npu-smi info -t spod-info -i {card_id} -c {chip_id} | grep \"Super Pod ID\""
|
|
).split(":")[1].strip()
|
|
else:
|
|
device_ip = get_cmd_stdout(
|
|
f"{hccn_tool_path} -i {device_id} -ip -g | grep ipaddr"
|
|
).split(":")[1].strip()
|
|
|
|
device_info = {
|
|
"server_id": local_host,
|
|
"device_id": str(device_id),
|
|
"device_ip": str(device_ip),
|
|
}
|
|
if soc_info == AscendSocVersion.A3:
|
|
device_info.update({
|
|
"super_pod_id": str(super_pod_id),
|
|
"super_device_id": str(super_device_id)
|
|
})
|
|
local_device_list.append(device_info)
|
|
|
|
dist.init_process_group(backend=dist.Backend.GLOO)
|
|
global_device_list = [None] * dist.get_world_size()
|
|
dist.all_gather_object(global_device_list, local_device_list)
|
|
global_device_list = [
|
|
device_info for device_list in global_device_list
|
|
for device_info in device_list # type: ignore[attr-defined]
|
|
]
|
|
cnt = 1
|
|
for device_info in global_device_list: # type: ignore[assignment]
|
|
device_info["cluster_id"] = str(cnt)
|
|
cnt += 1
|
|
assert (prefill_device_cnt + decode_device_cnt) <= len(global_device_list), \
|
|
"prefill_device_cnt + decode_device_cnt must be less than or equal to number of all devices in cluster"
|
|
ranktable = {
|
|
"version":
|
|
"1.2",
|
|
"server_count":
|
|
str(world_size),
|
|
"prefill_device_list":
|
|
global_device_list[:prefill_device_cnt],
|
|
"decode_device_list":
|
|
global_device_list[prefill_device_cnt:prefill_device_cnt +
|
|
decode_device_cnt],
|
|
"status":
|
|
"completed"
|
|
}
|
|
|
|
if local_rank == '0':
|
|
with open("ranktable.json", "w") as f:
|
|
json.dump(ranktable, f, indent=4)
|
|
|
|
print("gen ranktable.json done")
|