v0.10.1rc1
This commit is contained in:
246
examples/disaggregated_prefill_v1/README.md
Normal file
246
examples/disaggregated_prefill_v1/README.md
Normal file
@@ -0,0 +1,246 @@
|
||||
# Disaggregated Prefill-Decode Deployment Guide
|
||||
|
||||
## Overview
|
||||
This demo document provides instructions for running a disaggregated vLLM-ascend service with separate prefill and decode stages across 4 nodes, uses 16 Ascend NPUs for two prefill nodes (P1/P2) and 16 Ascend NPUS for two decode nodes (D1/D2).
|
||||
|
||||
## Prerequisites
|
||||
- Ascend NPU environment with vLLM 0.9.1 installed
|
||||
- Network interfaces configured for distributed communication (eg: eth0)
|
||||
- Model weights located at `/models/deepseek_r1_w8a8`
|
||||
|
||||
## Rank table generation
|
||||
The rank table is a JSON file that specifies the mapping of Ascend NPU ranks to nodes. The following command generates a rank table for all nodes with 16 cards prefill and 16 cards decode:
|
||||
|
||||
Run the following command on every node to generate the rank table:
|
||||
```shell
|
||||
cd /vllm-workspace/vllm-ascend/examples/disaggregated_prefill_v1/
|
||||
bash gen_ranktable.sh --ips 172.19.32.175 172.19.241.49 172.19.123.51 172.19.190.36 \
|
||||
--npus-per-node 8 --network-card-name eth0 --prefill-device-cnt 16 --decode-device-cnt 16
|
||||
```
|
||||
Rank table will generated at `/vllm-workspace/vllm-ascend/examples/disaggregated_prefill_v1/ranktable.json`
|
||||
|
||||
## Start disaggregated vLLM-ascend service
|
||||
For demonstration purposes, we will utilize the quantized version of Deepseek-R1. Recommended Parallelization Strategies:
|
||||
- P-node: DP2-TP8-EP16 (Data Parallelism 2, Tensor Parallelism 8, Expert Parallelism 16)
|
||||
- D-node: DP4-TP4-EP16 (Data Parallelism 4, Tensor Parallelism 4, Expert Parallelism 16)
|
||||
|
||||
Execution Sequence
|
||||
- 4 configured node ip are: 172.19.32.175 172.19.241.49 172.19.123.51 172.19.190.36
|
||||
- Start Prefill on Node 1 (P1)
|
||||
- Start Prefill on Node 2 (P2)
|
||||
- Start Decode on Node 1 (D1)
|
||||
- Start Decode on Node 2 (D2)
|
||||
- Start proxy server on Node1
|
||||
|
||||
Run prefill server P1 on first node:
|
||||
```shell
|
||||
export HCCL_IF_IP=172.19.32.175 # node ip
|
||||
export GLOO_SOCKET_IFNAME="eth0" # network card name
|
||||
export TP_SOCKET_IFNAME="eth0"
|
||||
export HCCL_SOCKET_IFNAME="eth0"
|
||||
export DISAGGREGATED_PREFILL_RANK_TABLE_PATH=/vllm-workspace/vllm-ascend/examples/disaggregated_prefill_v1/ranktable.json
|
||||
export OMP_PROC_BIND=false
|
||||
export OMP_NUM_THREADS=100
|
||||
export VLLM_USE_V1=1
|
||||
export VLLM_LLMDD_RPC_PORT=5559
|
||||
|
||||
vllm serve /models/deepseek_r1_w8a8 \
|
||||
--host 0.0.0.0 \
|
||||
--port 20002 \
|
||||
--data-parallel-size 2 \
|
||||
--data-parallel-size-local 1 \
|
||||
--api-server-count 2 \
|
||||
--data-parallel-address 172.19.32.175 \
|
||||
--data-parallel-rpc-port 13356 \
|
||||
--tensor-parallel-size 8 \
|
||||
--enable-expert-parallel \
|
||||
--seed 1024 \
|
||||
--served-model-name deepseek \
|
||||
--max-model-len 32768 \
|
||||
--max-num-batched-tokens 32768 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector": "LLMDataDistCMgrConnector",
|
||||
"kv_buffer_device": "npu",
|
||||
"kv_role": "kv_producer",
|
||||
"kv_parallel_size": 1,
|
||||
"kv_port": "20001",
|
||||
"engine_id": "0",
|
||||
"kv_connector_module_path": "vllm_ascend.distributed.llmdatadist_c_mgr_connector"
|
||||
}' \
|
||||
--additional-config \
|
||||
'{"chunked_prefill_for_mla":true}'
|
||||
```
|
||||
|
||||
Run prefill server P2 on second node:
|
||||
```shell
|
||||
export HCCL_IF_IP=172.19.241.49
|
||||
export GLOO_SOCKET_IFNAME="eth0"
|
||||
export TP_SOCKET_IFNAME="eth0"
|
||||
export HCCL_SOCKET_IFNAME="eth0"
|
||||
export DISAGGREGATED_PREFILL_RANK_TABLE_PATH=/vllm-workspace/vllm-ascend/examples/disaggregated_prefill_v1/ranktable.json
|
||||
export OMP_PROC_BIND=false
|
||||
export OMP_NUM_THREADS=100
|
||||
export VLLM_USE_V1=1
|
||||
export VLLM_LLMDD_RPC_PORT=5659
|
||||
|
||||
vllm serve /models/deepseek_r1_w8a8 \
|
||||
--host 0.0.0.0 \
|
||||
--port 20002 \
|
||||
--headless \
|
||||
--data-parallel-size 2 \
|
||||
--data-parallel-start-rank 1 \
|
||||
--data-parallel-size-local 1 \
|
||||
--data-parallel-address 172.19.32.175 \
|
||||
--data-parallel-rpc-port 13356 \
|
||||
--tensor-parallel-size 8 \
|
||||
--enable-expert-parallel \
|
||||
--seed 1024 \
|
||||
--served-model-name deepseek \
|
||||
--max-model-len 32768 \
|
||||
--max-num-batched-tokens 32768 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector": "LLMDataDistCMgrConnector",
|
||||
"kv_buffer_device": "npu",
|
||||
"kv_role": "kv_producer",
|
||||
"kv_parallel_size": 1,
|
||||
"kv_port": "20001",
|
||||
"engine_id": "0",
|
||||
"kv_connector_module_path": "vllm_ascend.distributed.llmdatadist_c_mgr_connector"
|
||||
}' \
|
||||
--additional-config \
|
||||
'{"chunked_prefill_for_mla":true}'
|
||||
```
|
||||
|
||||
Run decode server d1 on third node:
|
||||
|
||||
* In the D node, the `max-num-batched-tokens` parameter can be set to a smaller value since the D node processes at most `max-num-seqs` batches concurrently. As the `profile_run` only needs to handle `max-num-seqs` sequences at a time, we can safely set `max-num-batched-tokens` equal to `max-num-seqs`. This optimization will help reduce activation memory consumption.
|
||||
```shell
|
||||
export HCCL_IF_IP=172.19.123.51
|
||||
export GLOO_SOCKET_IFNAME="eth0"
|
||||
export TP_SOCKET_IFNAME="eth0"
|
||||
export HCCL_SOCKET_IFNAME="eth0"
|
||||
export DISAGGREGATED_PREFILL_RANK_TABLE_PATH=/vllm-workspace/vllm-ascend/examples/disaggregated_prefill_v1/ranktable.json
|
||||
export OMP_PROC_BIND=false
|
||||
export OMP_NUM_THREADS=100
|
||||
export VLLM_USE_V1=1
|
||||
export VLLM_LLMDD_RPC_PORT=5759
|
||||
|
||||
vllm serve /models/deepseek_r1_w8a8 \
|
||||
--host 0.0.0.0 \
|
||||
--port 20002 \
|
||||
--data-parallel-size 4 \
|
||||
--data-parallel-size-local 2 \
|
||||
--api-server-count 2 \
|
||||
--data-parallel-address 172.19.123.51 \
|
||||
--data-parallel-rpc-port 13356 \
|
||||
--tensor-parallel-size 4 \
|
||||
--enable-expert-parallel \
|
||||
--seed 1024 \
|
||||
--served-model-name deepseek \
|
||||
--max-model-len 32768 \
|
||||
--max-num-batched-tokens 256 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector": "LLMDataDistCMgrConnector",
|
||||
"kv_buffer_device": "npu",
|
||||
"kv_role": "kv_consumer",
|
||||
"kv_parallel_size": 1,
|
||||
"kv_port": "20001",
|
||||
"engine_id": "0",
|
||||
"kv_connector_module_path": "vllm_ascend.distributed.llmdatadist_c_mgr_connector"
|
||||
}' \
|
||||
--additional-config \
|
||||
'{"torchair_graph_config": {"enabled":true}}'
|
||||
```
|
||||
|
||||
Run decode server d2 on last node:
|
||||
```shell
|
||||
export HCCL_IF_IP=172.19.190.36
|
||||
export GLOO_SOCKET_IFNAME="eth0"
|
||||
export TP_SOCKET_IFNAME="eth0"
|
||||
export HCCL_SOCKET_IFNAME="eth0"
|
||||
export DISAGGREGATED_PREFILL_RANK_TABLE_PATH=/vllm-workspace/vllm-ascend/examples/disaggregated_prefill_v1/ranktable.json
|
||||
export OMP_PROC_BIND=false
|
||||
export OMP_NUM_THREADS=100
|
||||
export VLLM_USE_V1=1
|
||||
export VLLM_LLMDD_RPC_PORT=5859
|
||||
|
||||
vllm serve /models/deepseek_r1_w8a8 \
|
||||
--host 0.0.0.0 \
|
||||
--port 20002 \
|
||||
--headless \
|
||||
--data-parallel-size 4 \
|
||||
--data-parallel-start-rank 2 \
|
||||
--data-parallel-size-local 2 \
|
||||
--data-parallel-address 172.19.123.51 \
|
||||
--data-parallel-rpc-port 13356 \
|
||||
--tensor-parallel-size 4 \
|
||||
--enable-expert-parallel \
|
||||
--seed 1024 \
|
||||
--served-model-name deepseek \
|
||||
--max-model-len 32768 \
|
||||
--max-num-batched-tokens 256 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector": "LLMDataDistCMgrConnector",
|
||||
"kv_buffer_device": "npu",
|
||||
"kv_role": "kv_consumer",
|
||||
"kv_parallel_size": 1,
|
||||
"kv_port": "20001",
|
||||
"engine_id": "0",
|
||||
"kv_connector_module_path": "vllm_ascend.distributed.llmdatadist_c_mgr_connector"
|
||||
}' \
|
||||
--additional-config \
|
||||
'{"torchair_graph_config": {"enabled":true}}'
|
||||
```
|
||||
|
||||
Run proxy server on the first node:
|
||||
```shell
|
||||
cd /vllm-workspace/vllm-ascend/examples/disaggregated_prefill_v1
|
||||
python toy_proxy_server.py --host 172.19.32.175 --port 1025 --prefiller-hosts 172.19.241.49 --prefiller-port 20002 --decoder-hosts 172.19.123.51 --decoder-ports 20002
|
||||
```
|
||||
|
||||
Verification
|
||||
Check service health using the proxy server endpoint:
|
||||
```shell
|
||||
curl http://localhost:1025/v1/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "deepseek",
|
||||
"prompt": "Who are you?",
|
||||
"max_tokens": 100,
|
||||
"temperature": 0
|
||||
}'
|
||||
```
|
||||
|
||||
Performance
|
||||
Test performance with vllm benchmark:
|
||||
```shell
|
||||
cd /vllm-workspace/vllm/benchmarks
|
||||
python3 benchmark_serving.py \
|
||||
--backend vllm \
|
||||
--dataset-name random \
|
||||
--random-input-len 4096 \
|
||||
--random-output-len 1536 \
|
||||
--num-prompts 256 \
|
||||
--ignore-eos \
|
||||
--model deepseek \
|
||||
--tokenizer /models/deepseek_r1_w8a8 \
|
||||
--host localhost \
|
||||
--port 1025 \
|
||||
--endpoint /v1/completions \
|
||||
--max-concurrency 4 \
|
||||
--request-rate 4
|
||||
```
|
||||
122
examples/disaggregated_prefill_v1/gen_ranktable.py
Normal file
122
examples/disaggregated_prefill_v1/gen_ranktable.py
Normal file
@@ -0,0 +1,122 @@
|
||||
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")
|
||||
79
examples/disaggregated_prefill_v1/gen_ranktable.sh
Normal file
79
examples/disaggregated_prefill_v1/gen_ranktable.sh
Normal file
@@ -0,0 +1,79 @@
|
||||
#!/bin/bash
|
||||
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh
|
||||
export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/op_api/lib/:${LD_LIBRARY_PATH}
|
||||
|
||||
NPUS_PER_NODE=8
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--ips)
|
||||
shift
|
||||
while [[ $# -gt 0 && ! "$1" == --* ]]; do
|
||||
IPs+=("$1")
|
||||
shift
|
||||
done
|
||||
;;
|
||||
--npus-per-node)
|
||||
shift
|
||||
NPUS_PER_NODE="$1"
|
||||
shift
|
||||
;;
|
||||
--network-card-name)
|
||||
shift
|
||||
NETWORK_CARD_NAME="$1"
|
||||
shift
|
||||
;;
|
||||
--prefill-device-cnt)
|
||||
shift
|
||||
PREFILL_DEVICE_CNT="$1"
|
||||
shift
|
||||
;;
|
||||
--decode-device-cnt)
|
||||
shift
|
||||
DECODE_DEVICE_CNT="$1"
|
||||
shift
|
||||
;;
|
||||
esac
|
||||
done
|
||||
LOCAL_HOSTS=($(hostname -I))
|
||||
LOCAL_HOST="127.0.0.1"
|
||||
MASTER_ADDR=${IPs[0]}
|
||||
MASTER_PORT=6657
|
||||
NNODES=${#IPs[@]}
|
||||
NODE_RANK="8"
|
||||
for i in "${!IPs[@]}"; do
|
||||
ip="${IPs[$i]}"
|
||||
for local_host in "${LOCAL_HOSTS[@]}"; do
|
||||
if [[ "$local_host" == "$ip" ]]; then
|
||||
LOCAL_HOST=$local_host
|
||||
NODE_RANK=$i
|
||||
break 2
|
||||
fi
|
||||
done
|
||||
done
|
||||
|
||||
if [[ $NODE_RANK == "" ]];then
|
||||
echo "[Error] para \"NODE_RANK\" must be defined"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
WORLD_SIZE=$(($NPUS_PER_NODE * $NNODES))
|
||||
RANKSTART=`expr $NPUS_PER_NODE \* $NODE_RANK`
|
||||
|
||||
echo "========>param:"
|
||||
echo "LOCAL_HOST": $LOCAL_HOST
|
||||
echo "WORLD_SIZE: " $WORLD_SIZE
|
||||
echo "RANKSTART": $RANKSTART
|
||||
echo "NNODES": $NNODES
|
||||
echo "NODE_RANK": $NODE_RANK
|
||||
echo "==============="
|
||||
|
||||
if [[ -n "${GEN_RANKTABLE}" || ! -e ${PWD}/ranktable.json ]]; then
|
||||
GLOO_SOCKET_IFNAME=$NETWORK_CARD_NAME torchrun \
|
||||
--nproc_per_node 1 \
|
||||
--nnodes ${NNODES} \
|
||||
--node_rank ${NODE_RANK} \
|
||||
--master_addr ${MASTER_ADDR} \
|
||||
--master_port ${MASTER_PORT} \
|
||||
gen_ranktable.py --local-host $LOCAL_HOST --prefill-device-cnt $PREFILL_DEVICE_CNT --decode-device-cnt $DECODE_DEVICE_CNT
|
||||
fi
|
||||
@@ -0,0 +1,546 @@
|
||||
# Adapted from https://github.com/vllm-project/vllm/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Tutorial: Using the Load Balance Proxy Server Example
|
||||
#
|
||||
# This proxy server is designed to distribute requests between multiple
|
||||
# "prefiller" and "decoder" backend servers for large language model inference.
|
||||
# It is useful for scaling out inference workloads and balancing load across
|
||||
# multiple backend instances.
|
||||
#
|
||||
# Features:
|
||||
# - Load balances requests to multiple prefiller and decoder servers.
|
||||
# - Supports OpenAI-compatible /v1/completions and /v1/chat/completions endpoints.
|
||||
# - Streams responses from backend servers to clients.
|
||||
#
|
||||
# Prerequisites:
|
||||
# - Python 3.8+
|
||||
# - Install dependencies:
|
||||
# pip install fastapi httpx uvicorn vllm
|
||||
#
|
||||
# Step 1: Start Your Backend Servers
|
||||
# ----------------------------------
|
||||
# You need to have at least one prefiller and one decoder backend running.
|
||||
# These can be mock servers or actual vLLM servers.
|
||||
#
|
||||
# For testing, you can use the provided mock server:
|
||||
#
|
||||
# vllm serve --host 0.0.0.0 --port 8100 ... # Prefiller 1
|
||||
# vllm serve --host 0.0.0.0 --port 8101 ... # Prefiller 2
|
||||
# vllm serve --host 0.0.0.0 --port 8200 ... # Decoder 1
|
||||
# vllm serve --host 0.0.0.0 --port 8201 ... # Decoder 2
|
||||
#
|
||||
# Step 2: Start the Proxy Server
|
||||
# ------------------------------
|
||||
# Run the proxy server, specifying the host/port for each prefiller and decoder:
|
||||
#
|
||||
# python load_balance_proxy_server_example.py \
|
||||
# --host 0.0.0.0 --port 9000 \
|
||||
# --prefiller-hosts 127.0.0.1 127.0.0.1 \
|
||||
# --prefiller-ports 8100 8101 \
|
||||
# --decoder-hosts 127.0.0.1 127.0.0.1 \
|
||||
# --decoder-ports 8200 8201
|
||||
#
|
||||
# This will start the proxy on port 9000, load balancing between two prefiller
|
||||
# and two decoder servers.
|
||||
#
|
||||
# Step 3: Send a Request to the Proxy
|
||||
# -----------------------------------
|
||||
# You can now send OpenAI-compatible requests to the proxy. For example:
|
||||
#
|
||||
# curl -X POST http://localhost:9000/v1/completions \
|
||||
# -H "Content-Type: application/json" \
|
||||
# -d '{
|
||||
# "model": "your-model",
|
||||
# "prompt": "The quick brown fox jumps over the lazy dog",
|
||||
# "max_tokens": 16
|
||||
# }'
|
||||
#
|
||||
# Or for chat completions:
|
||||
#
|
||||
# curl -X POST http://localhost:9000/v1/chat/completions \
|
||||
# -H "Content-Type: application/json" \
|
||||
# -d '{
|
||||
# "model": "your-model",
|
||||
# "messages": [{"role": "user", "content": "Hello!"}],
|
||||
# "max_tokens": 16
|
||||
# }'
|
||||
#
|
||||
# Step 4: Health Check
|
||||
# --------------------
|
||||
# To check if the proxy is running and see how many backend instances are
|
||||
# connected, use:
|
||||
#
|
||||
# curl http://localhost:9000/healthcheck
|
||||
#
|
||||
# This will return a JSON object with the status and the number of prefiller
|
||||
# and decoder instances.
|
||||
#
|
||||
# Notes:
|
||||
# - You can scale the number of prefiller and decoder servers as needed.
|
||||
# - The proxy will round-robin requests to balance load.
|
||||
# - For production, ensure your backend servers are robust and secure.
|
||||
#
|
||||
# For more details, see the code and comments in this file.
|
||||
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import functools
|
||||
import heapq
|
||||
import os
|
||||
import sys
|
||||
import uuid
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import List
|
||||
|
||||
import httpx
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
from vllm.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Add uvloop for faster event loop if available
|
||||
try:
|
||||
import uvloop
|
||||
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
class ServerState:
|
||||
|
||||
def __init__(self, host, port):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.url = f'http://{host}:{port}/v1'
|
||||
self.client = httpx.AsyncClient(timeout=None,
|
||||
base_url=self.url,
|
||||
limits=httpx.Limits(
|
||||
max_connections=100000,
|
||||
max_keepalive_connections=100000))
|
||||
self.active_tokens = 0
|
||||
self.active_kv_cache = 0 # Only for prefiller
|
||||
self.active_requests = 0 # Number of active requests
|
||||
self.aborted_requests = set() # Track aborted requests
|
||||
# Removed individual server lock - will use global locks instead
|
||||
|
||||
|
||||
class ProxyState:
|
||||
|
||||
def __init__(self, prefiller_instances, decoder_instances):
|
||||
self.prefillers: List[ServerState] = [
|
||||
ServerState(h, p) for h, p in prefiller_instances
|
||||
]
|
||||
self.decoders: List[ServerState] = [
|
||||
ServerState(h, p) for h, p in decoder_instances
|
||||
]
|
||||
self.req_to_prefiller = {}
|
||||
self.req_id_lock = asyncio.Lock()
|
||||
# Removed selection locks - no longer needed for synchronous methods
|
||||
|
||||
# Initialize priority queues for efficient server selection
|
||||
# Each entry is (priority_score, server_index, server_reference)
|
||||
# Lower priority score = higher priority (less loaded)
|
||||
self.prefiller_heap = [(0, i, server)
|
||||
for i, server in enumerate(self.prefillers)]
|
||||
self.decoder_heap = [(0, i, server)
|
||||
for i, server in enumerate(self.decoders)]
|
||||
heapq.heapify(self.prefiller_heap)
|
||||
heapq.heapify(self.decoder_heap)
|
||||
|
||||
def _update_prefiller_priority(self, server_idx: int):
|
||||
"""Update the priority of a prefiller server in the heap."""
|
||||
server = self.prefillers[server_idx]
|
||||
# Priority based on active_tokens and active_kv_cache
|
||||
priority = server.active_tokens + server.active_kv_cache * 0.3
|
||||
# Remove old entry and add new one
|
||||
self.prefiller_heap = [(p, i, s) for p, i, s in self.prefiller_heap
|
||||
if i != server_idx]
|
||||
heapq.heappush(self.prefiller_heap,
|
||||
(priority, server_idx, server)) # type: ignore
|
||||
|
||||
def _update_decoder_priority(self, server_idx: int):
|
||||
"""Update the priority of a decoder server in the heap."""
|
||||
server = self.decoders[server_idx]
|
||||
priority = server.active_tokens
|
||||
# Remove old entry and add new one
|
||||
self.decoder_heap = [(p, i, s) for p, i, s in self.decoder_heap
|
||||
if i != server_idx]
|
||||
heapq.heappush(self.decoder_heap,
|
||||
(priority, server_idx, server)) # type: ignore
|
||||
|
||||
def abort_prefiller_request(self, server_idx: int,
|
||||
request_id): # Changed to synchronous
|
||||
"""
|
||||
Mark a request as aborted. This will helps to release kv cache in
|
||||
prefiller node.
|
||||
"""
|
||||
# No lock needed - atomic operation
|
||||
self.prefillers[server_idx].aborted_requests.add(request_id)
|
||||
|
||||
def aquire_aborted_prefiller_requests(
|
||||
self, server_idx: int): # Changed to synchronous
|
||||
"""
|
||||
Get the set of aborted requests and clear it.
|
||||
This is used to release kv cache in prefiller node.
|
||||
"""
|
||||
# No lock needed - atomic operation
|
||||
aborted_requests = self.prefillers[server_idx].aborted_requests.copy()
|
||||
self.prefillers[server_idx].aborted_requests.clear()
|
||||
return aborted_requests
|
||||
|
||||
async def next_req_id(self):
|
||||
async with self.req_id_lock:
|
||||
return str(uuid.uuid4())
|
||||
|
||||
def select_prefiller(self, token_count): # Changed to synchronous
|
||||
# No lock needed - entire function is atomic
|
||||
if not self.prefiller_heap:
|
||||
raise RuntimeError("No prefiller servers available")
|
||||
|
||||
priority, chosen, server = heapq.heappop(self.prefiller_heap)
|
||||
|
||||
# Update the chosen server atomically
|
||||
self.prefillers[chosen].active_tokens += token_count
|
||||
self.prefillers[chosen].active_kv_cache += token_count
|
||||
|
||||
# Update priority and re-add to heap
|
||||
self._update_prefiller_priority(chosen)
|
||||
|
||||
return chosen
|
||||
|
||||
def release_prefiller(self, idx, token_count): # Changed to synchronous
|
||||
# No lock needed - atomic operation
|
||||
self.prefillers[idx].active_tokens -= token_count
|
||||
# Update priority queue after releasing
|
||||
self._update_prefiller_priority(idx)
|
||||
|
||||
def release_prefiller_kv(self, idx, token_count): # Changed to synchronous
|
||||
# No lock needed - atomic operation
|
||||
if self.prefillers[idx].active_kv_cache > 0:
|
||||
self.prefillers[idx].active_kv_cache -= token_count
|
||||
# Update priority queue after releasing
|
||||
self._update_prefiller_priority(idx)
|
||||
|
||||
def select_decoder(self, token_count): # Changed to synchronous
|
||||
# No lock needed - entire function is atomic
|
||||
if not self.decoder_heap:
|
||||
raise RuntimeError("No decoder servers available")
|
||||
|
||||
priority, chosen, server = heapq.heappop(self.decoder_heap)
|
||||
|
||||
# Update the chosen server atomically
|
||||
self.decoders[chosen].active_tokens += token_count
|
||||
|
||||
# Update priority and re-add to heap
|
||||
self._update_decoder_priority(chosen)
|
||||
|
||||
return chosen
|
||||
|
||||
def release_decoder(self, idx, token_count): # Changed to synchronous
|
||||
# No lock needed - atomic operation
|
||||
self.decoders[idx].active_tokens -= token_count
|
||||
# Update priority queue after releasing
|
||||
self._update_decoder_priority(idx)
|
||||
|
||||
# Omni_infer's calculate_input_scores function
|
||||
def calculate_prefill_scores(self, request_length: int) -> float:
|
||||
length_score = request_length / 4.0
|
||||
input_score = length_score * 0.0345 + 120.0745
|
||||
return input_score
|
||||
|
||||
def calculate_decode_scores(self, request_length: int) -> float:
|
||||
return request_length
|
||||
|
||||
|
||||
proxy_state = None
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--prefiller-hosts",
|
||||
type=str,
|
||||
nargs="+",
|
||||
default=["localhost"])
|
||||
parser.add_argument("--prefiller-ports",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[8001])
|
||||
parser.add_argument("--decoder-hosts",
|
||||
type=str,
|
||||
nargs="+",
|
||||
default=["localhost"])
|
||||
parser.add_argument("--decoder-ports", type=int, nargs="+", default=[8002])
|
||||
parser.add_argument("--max-retries",
|
||||
type=int,
|
||||
default=3,
|
||||
help="Maximum number of retries for HTTP requests")
|
||||
parser.add_argument(
|
||||
"--retry-delay",
|
||||
type=float,
|
||||
default=0.001,
|
||||
help="Base delay (seconds) for exponential backoff retries")
|
||||
args = parser.parse_args()
|
||||
if len(args.prefiller_hosts) != len(args.prefiller_ports):
|
||||
raise ValueError(
|
||||
"Number of prefiller hosts must match number of prefiller ports")
|
||||
if len(args.decoder_hosts) != len(args.decoder_ports):
|
||||
raise ValueError(
|
||||
"Number of decoder hosts must match number of decoder ports")
|
||||
args.prefiller_instances = list(
|
||||
zip(args.prefiller_hosts, args.prefiller_ports))
|
||||
args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports))
|
||||
return args
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
global proxy_state
|
||||
proxy_state = ProxyState(global_args.prefiller_instances,
|
||||
global_args.decoder_instances)
|
||||
print(
|
||||
f"Initialized {len(proxy_state.prefillers)} prefill clients and {len(proxy_state.decoders)} decode clients."
|
||||
)
|
||||
yield
|
||||
for p in proxy_state.prefillers:
|
||||
await p.client.aclose()
|
||||
for d in proxy_state.decoders:
|
||||
await d.client.aclose()
|
||||
|
||||
|
||||
async def listen_for_disconnect(request: Request) -> None:
|
||||
"""Return if a disconnect message is received"""
|
||||
while True:
|
||||
message = await request.receive()
|
||||
if message["type"] == "http.disconnect":
|
||||
break
|
||||
|
||||
|
||||
def with_cancellation(handler_func):
|
||||
|
||||
@functools.wraps(handler_func)
|
||||
async def wrapper(*args, **kwargs):
|
||||
request = kwargs["request"]
|
||||
handler_task = asyncio.create_task(handler_func(*args, **kwargs))
|
||||
cancellation_task = asyncio.create_task(listen_for_disconnect(request))
|
||||
done, pending = await asyncio.wait([handler_task, cancellation_task],
|
||||
return_when=asyncio.FIRST_COMPLETED)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
if handler_task in done:
|
||||
return handler_task.result()
|
||||
return None
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
app = FastAPI(lifespan=lifespan)
|
||||
|
||||
|
||||
async def send_request_to_service(client: httpx.AsyncClient,
|
||||
prefiller_id: int,
|
||||
endpoint: str,
|
||||
req_data: dict,
|
||||
request_id: str,
|
||||
max_retries: int = 3,
|
||||
base_delay: float = 0.2):
|
||||
aborted_requests = proxy_state.aquire_aborted_prefiller_requests(
|
||||
prefiller_id)
|
||||
req_data = req_data.copy()
|
||||
req_data['kv_transfer_params'] = {
|
||||
"do_remote_decode": True,
|
||||
"do_remote_prefill": False,
|
||||
"remote_engine_id": None,
|
||||
"remote_block_ids": None,
|
||||
"remote_host": None,
|
||||
"remote_port": None,
|
||||
"aborted_request": list(aborted_requests),
|
||||
}
|
||||
req_data["stream"] = False
|
||||
req_data["max_tokens"] = 1
|
||||
if "stream_options" in req_data:
|
||||
del req_data["stream_options"]
|
||||
headers = {
|
||||
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
|
||||
"X-Request-Id": request_id
|
||||
}
|
||||
last_exc = None
|
||||
for attempt in range(1, max_retries + 1):
|
||||
try:
|
||||
response = await client.post(endpoint,
|
||||
json=req_data,
|
||||
headers=headers)
|
||||
response.raise_for_status()
|
||||
return response
|
||||
except (httpx.RequestError, httpx.HTTPStatusError) as e:
|
||||
logger.warning(
|
||||
f"Attempt {attempt} failed for {endpoint}: {str(e)}")
|
||||
last_exc = e
|
||||
if attempt < max_retries:
|
||||
await asyncio.sleep(base_delay * (2**(attempt - 1)))
|
||||
else:
|
||||
logger.error(
|
||||
f"All {max_retries} attempts failed for {endpoint}.")
|
||||
raise last_exc
|
||||
|
||||
|
||||
async def stream_service_response_with_retry(client: httpx.AsyncClient,
|
||||
endpoint: str,
|
||||
req_data: dict,
|
||||
request_id: str,
|
||||
max_retries: int = 3,
|
||||
base_delay: float = 0.2):
|
||||
headers = {
|
||||
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
|
||||
"X-Request-Id": request_id
|
||||
}
|
||||
for attempt in range(1, max_retries + 1):
|
||||
try:
|
||||
async with client.stream("POST",
|
||||
endpoint,
|
||||
json=req_data,
|
||||
headers=headers) as response:
|
||||
response.raise_for_status()
|
||||
first_chunk_sent = False
|
||||
async for chunk in response.aiter_bytes():
|
||||
first_chunk_sent = True
|
||||
yield chunk
|
||||
return # Success, exit after streaming
|
||||
except (httpx.RequestError, httpx.HTTPStatusError) as e:
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"Attempt {attempt} failed for streaming {endpoint}: {str(e)}"
|
||||
)
|
||||
await asyncio.sleep(base_delay * (2**(attempt - 1)))
|
||||
else:
|
||||
logger.error(
|
||||
f"All {max_retries} attempts failed for streaming {endpoint}."
|
||||
)
|
||||
raise e
|
||||
except Exception as e:
|
||||
# If any chunk has been sent, do not retry, just log and drop
|
||||
if 'first_chunk_sent' in locals() and first_chunk_sent:
|
||||
logger.error(
|
||||
f"Streaming to client interrupted after response started: {str(e)}"
|
||||
)
|
||||
return
|
||||
else:
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"Attempt {attempt} failed for streaming {endpoint}: {str(e)}"
|
||||
)
|
||||
await asyncio.sleep(base_delay * (2**(attempt - 1)))
|
||||
else:
|
||||
logger.error(
|
||||
f"All {max_retries} attempts failed for streaming {endpoint}."
|
||||
)
|
||||
raise e
|
||||
|
||||
|
||||
async def _handle_completions(api: str, request: Request):
|
||||
try:
|
||||
req_data = await request.json()
|
||||
req_body = await request.body()
|
||||
request_length = len(req_body)
|
||||
prefiller_score = proxy_state.calculate_prefill_scores(request_length)
|
||||
logger.debug(
|
||||
f"Request length: {request_length}, Prefiller score: {prefiller_score}"
|
||||
)
|
||||
request_id = await proxy_state.next_req_id()
|
||||
# Select prefiller
|
||||
prefiller_idx = proxy_state.select_prefiller(prefiller_score)
|
||||
prefiller = proxy_state.prefillers[prefiller_idx]
|
||||
# Send request to prefiller
|
||||
response = await send_request_to_service(
|
||||
prefiller.client,
|
||||
prefiller_idx,
|
||||
api,
|
||||
req_data,
|
||||
request_id,
|
||||
max_retries=global_args.max_retries,
|
||||
base_delay=global_args.retry_delay)
|
||||
proxy_state.release_prefiller(prefiller_idx, prefiller_score)
|
||||
response_json = response.json()
|
||||
kv_transfer_params = response_json.get('kv_transfer_params', {})
|
||||
if kv_transfer_params:
|
||||
req_data["kv_transfer_params"] = kv_transfer_params
|
||||
# Select decoder
|
||||
decoder_score = proxy_state.calculate_decode_scores(request_length)
|
||||
logger.debug("Decoder score: %f", decoder_score)
|
||||
# Use the prefiller's kv_transfer_params to select decoder
|
||||
decoder_idx = proxy_state.select_decoder(decoder_score)
|
||||
decoder = proxy_state.decoders[decoder_idx]
|
||||
logger.debug("Using %s %s", prefiller.url, decoder.url)
|
||||
# Stream response from decoder
|
||||
released_kv = False
|
||||
|
||||
async def generate_stream():
|
||||
nonlocal released_kv
|
||||
# Only one await per chunk, minimal logic in loop
|
||||
try:
|
||||
async for chunk in stream_service_response_with_retry(
|
||||
decoder.client,
|
||||
api,
|
||||
req_data,
|
||||
request_id=request_id,
|
||||
max_retries=global_args.max_retries,
|
||||
base_delay=global_args.retry_delay):
|
||||
if not released_kv and chunk:
|
||||
proxy_state.release_prefiller_kv(
|
||||
prefiller_idx, prefiller_score)
|
||||
released_kv = True
|
||||
yield chunk
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error during streaming from decoder {decoder.url}: {str(e)} the aborted request {request_id} will be routing to the target prefiller when new request is ready to dispatch to it"
|
||||
)
|
||||
proxy_state.abort_prefiller_request(prefiller_idx, request_id)
|
||||
proxy_state.release_prefiller_kv(prefiller_idx,
|
||||
prefiller_score)
|
||||
|
||||
# After streaming done, release tokens
|
||||
proxy_state.release_decoder(decoder_idx, decoder_score)
|
||||
|
||||
return StreamingResponse(generate_stream(),
|
||||
media_type="application/json")
|
||||
except Exception as e:
|
||||
import traceback
|
||||
exc_info = sys.exc_info()
|
||||
print("Error occurred in disagg prefill proxy server"
|
||||
f" - {api} endpoint")
|
||||
print(e)
|
||||
print("".join(traceback.format_exception(*exc_info)))
|
||||
raise
|
||||
|
||||
|
||||
@app.post("/v1/completions")
|
||||
@with_cancellation
|
||||
async def handle_completions(request: Request):
|
||||
return await _handle_completions("/completions", request)
|
||||
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
@with_cancellation
|
||||
async def handle_chat_completions(request: Request):
|
||||
return await _handle_completions("/chat/completions", request)
|
||||
|
||||
|
||||
@app.get("/healthcheck")
|
||||
async def healthcheck():
|
||||
return {
|
||||
"status": "ok",
|
||||
"prefill_instances": len(proxy_state.prefillers),
|
||||
"decode_instances": len(proxy_state.decoders)
|
||||
}
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
global global_args
|
||||
global_args = parse_args()
|
||||
import uvicorn
|
||||
uvicorn.run(app, host=global_args.host, port=global_args.port)
|
||||
@@ -0,0 +1,165 @@
|
||||
# Mooncake connector deployment Guide
|
||||
|
||||
## Environmental Dependencies
|
||||
|
||||
* Software:
|
||||
* Python >= 3.9, < 3.12
|
||||
* CANN >= 8.2.rc1
|
||||
* PyTorch >= 2.7.1, torch-npu >= 2.7.1.dev20250724
|
||||
* vLLM (same version as vllm-ascend)
|
||||
* mooncake-transfer-engine reference documentation: https://github.com/kvcache-ai/Mooncake/blob/main/doc/zh/ascend_transport.md
|
||||
|
||||
The vllm version must be the same as the main branch of vllm-ascend, for example, 2025/07/30. The version is
|
||||
|
||||
* vllm: v0.10.1
|
||||
* vllm-ascend: v0.10.1rc1
|
||||
|
||||
## run
|
||||
|
||||
### 1.Run `prefill` Node
|
||||
|
||||
```
|
||||
bash run_prefill.sh
|
||||
```
|
||||
|
||||
Content of the run_prefill.sh script
|
||||
|
||||
```
|
||||
export HCCL_EXEC_TIMEOUT=204
|
||||
export HCCL_CONNECT_TIMEOUT=120
|
||||
export HCCL_IF_IP=localhost
|
||||
export GLOO_SOCKET_IFNAME="xxxxxx"
|
||||
export TP_SOCKET_IFNAME="xxxxxx"
|
||||
export HCCL_SOCKET_IFNAME="xxxxxx"
|
||||
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3
|
||||
export PHYSICAL_DEVICES=$(ls /dev/davinci* 2>/dev/null | grep -o '[0-9]\+' | sort -n | paste -sd',' -)
|
||||
|
||||
vllm serve "/xxxxx/DeepSeek-V2-Lite-Chat" \
|
||||
--host localhost \
|
||||
--port 8100 \
|
||||
--tensor-parallel-size 2\
|
||||
--seed 1024 \
|
||||
--max-model-len 2000 \
|
||||
--max-num-batched-tokens 2000 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--data-parallel-size 2 \
|
||||
--data-parallel-address localhost \
|
||||
--data-parallel-rpc-port 9100 \
|
||||
--gpu-memory-utilization 0.8 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector": "MooncakeConnectorV1",
|
||||
"kv_buffer_device": "npu",
|
||||
"kv_role": "kv_producer",
|
||||
"kv_parallel_size": 1,
|
||||
"kv_port": "20001",
|
||||
"engine_id": "0",
|
||||
"kv_rank": 0,
|
||||
"kv_connector_module_path": "vllm_ascend.distributed.mooncake_connector",
|
||||
"kv_connector_extra_config": {
|
||||
"prefill": {
|
||||
"dp_size": 2,
|
||||
"tp_size": 2
|
||||
},
|
||||
"decode": {
|
||||
"dp_size": 2,
|
||||
"tp_size": 2
|
||||
}
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
`HCCL_EXEC_TIMEOUT`, `HCCL_CONNECT_TIMEOUT`, and `HCCL_IF_IP` are hccl-related configurations.<br>
|
||||
Set `GLOO_SOCKET_IFNAME`, `TP_SOCKET_IFNAME`, and `HCCL_SOCKET_IFNAME` to the corresponding NIC.<br>
|
||||
`ASCEND_RT_VISIBLE_DEVICES` specifies the cards on which the node run resides. The total number of cards equals `dp_size*tp_size`.<br>
|
||||
`/xxxxx/DeepSeek-V2-Lite-Chat` is configured as a model that requires run.<br>
|
||||
`--host`: indicates the IP address of the node to be started.<br>
|
||||
`--port`: indicates the port to be started, which corresponds to the port in step 4.<br>
|
||||
`--seed`, --max-model-len, and --max-num-batched-tokens model basic configuration. Set this parameter based on the site requirements.<br>
|
||||
`--tensor-parallel-size`: specifies the TP size.<br>
|
||||
`--data-parallel-size`: indicates the DP size.<br>
|
||||
`--data-parallel-address`: indicates the IP address of the DP. Set this parameter to the IP address of the node.--data-parallel-rpc-port: indicates the RPC port for communication in the DP group.<br>
|
||||
`--trust-remote-code` can load the local model.<br>
|
||||
`--enforce-eager` Turn off the map mode<br>
|
||||
`--gpu-memory-utilization`: Percentage of video memory occupied by the card<br>
|
||||
`--kv-transfer-config`: follow kv_connector, kv_connector_module_path: mooncakeconnect, kv_buffer_device, and run on the NPU card. For kv_role, set kv_producer to the p node, kv_consumer to the d node, kv_parallel_size to 1, and kv_port to the port used by the node. For the p node, set engine_id and kv_rank to 0 and for the d node to 1. Configure the distributed parallel policy for the p and d nodes in the kv_connector_extra_config file based on --tensor-parallel-size and --data-parallel-size.<br>
|
||||
|
||||
|
||||
### 2. Run `decode` Node
|
||||
|
||||
```
|
||||
bash run_decode.sh
|
||||
```
|
||||
|
||||
Content of the run_decode.sh script
|
||||
|
||||
```
|
||||
export HCCL_EXEC_TIMEOUT=204
|
||||
export HCCL_CONNECT_TIMEOUT=120
|
||||
export HCCL_IF_IP=localhost
|
||||
export GLOO_SOCKET_IFNAME="xxxxxx"
|
||||
export TP_SOCKET_IFNAME="xxxxxx"
|
||||
export HCCL_SOCKET_IFNAME="xxxxxx"
|
||||
export ASCEND_RT_VISIBLE_DEVICES=4,5,6,7
|
||||
export PHYSICAL_DEVICES=$(ls /dev/davinci* 2>/dev/null | grep -o '[0-9]\+' | sort -n | paste -sd',' -)
|
||||
|
||||
vllm serve "/xxxxx/DeepSeek-V2-Lite-Chat" \
|
||||
--host localhost \
|
||||
--port 8200 \
|
||||
--tensor-parallel-size 2\
|
||||
--seed 1024 \
|
||||
--max-model-len 2000 \
|
||||
--max-num-batched-tokens 2000 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--data-parallel-size 2 \
|
||||
--data-parallel-address localhost \
|
||||
--data-parallel-rpc-port 9100 \
|
||||
--gpu-memory-utilization 0.8 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector": "MooncakeConnectorV1",
|
||||
"kv_buffer_device": "npu",
|
||||
"kv_role": "kv_consumer",
|
||||
"kv_parallel_size": 1,
|
||||
"kv_port": "20002",
|
||||
"engine_id": "1",
|
||||
"kv_rank": 1,
|
||||
"kv_connector_module_path": "vllm_ascend.distributed.mooncake_connector",
|
||||
"kv_connector_extra_config": {
|
||||
"prefill": {
|
||||
"dp_size": 2,
|
||||
"tp_size": 2
|
||||
},
|
||||
"decode": {
|
||||
"dp_size": 2,
|
||||
"tp_size": 2
|
||||
}
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
### 3. Start proxy_server. ###
|
||||
|
||||
```
|
||||
cd /vllm-ascend/examples/disaggregate_prefill_v1/
|
||||
python load_balance_proxy_server_example.py --host localhost --prefiller-hosts host1 host2 --prefiller-ports 8100 8101 --decoder-hosts host3 host4 --decoder-ports 8200 8201
|
||||
```
|
||||
|
||||
`--host`: indicates the active node. The value of localhost in the curl command delivered in step 5 must be the same as the host. The default port number for starting the service proxy is 8000.<br>
|
||||
`--prefiller-hosts`: Set this parameter to the IP addresses of all p nodes. In the xpyd scenario, add the IP addresses to the end of this configuration item and leave a blank space between the IP addresses.<br>
|
||||
`--prefiller-ports`: Set this parameter to the port number of all p nodes, which is the configuration of the port number for the vllm to start the service in step 3. Write the port number after the configuration in sequence and leave a blank space between the port number and the port number. The sequence must be one-to-one mapping to the IP address of --prefiller-hosts.<br>
|
||||
`--decoder-hosts`: Set this parameter to the IP addresses of all d nodes. In the xpyd scenario, add the IP addresses to the end of this configuration item and leave a blank space between the IP addresses.<br>
|
||||
`--decoder-ports`: Set this parameter to the port number of all d nodes, which is the configuration of the port number for the vllm to start the service in step 4. Set port to the end of the configuration, and leave a blank space between port and port. The sequence must be one-to-one mapping to the IP address of --decoder-hosts.<br>
|
||||
|
||||
|
||||
### 4. Run Inference
|
||||
|
||||
Set the IP address in the inference file to the actual IP address. Set the model variable to the path of the model. Ensure that the path is the same as that in the shell script.
|
||||
|
||||
```
|
||||
curl -s http://localhost:8000/v1/completions -H "Content-Type: application/json" -d '{
|
||||
"model": "model_path",
|
||||
"prompt": "Given the accelerating impacts of climate change—including rising sea levels, increasing frequency of extreme weather events, loss of biodiversity, and adverse effects on agriculture and human health—there is an urgent need for a robust, globally coordinated response. However, international efforts are complicated by a range of factors: economic disparities between high-income and low-income countries, differing levels of industrialization, varying access to clean energy technologies, and divergent political systems that influence climate policy implementation. In this context, how can global agreements like the Paris Accord be redesigned or strengthened to not only encourage but effectively enforce emission reduction targets? Furthermore, what mechanisms can be introduced to promote fair and transparent technology transfer, provide adequate financial support for climate adaptation in vulnerable regions, and hold nations accountable without exacerbating existing geopolitical tensions or disproportionately burdening those with historically lower emissions?",
|
||||
"max_tokens": 256
|
||||
}'
|
||||
```
|
||||
32
examples/disaggregated_prefill_v1/run_server.sh
Normal file
32
examples/disaggregated_prefill_v1/run_server.sh
Normal file
@@ -0,0 +1,32 @@
|
||||
export HCCL_IF_IP=141.61.39.117
|
||||
export GLOO_SOCKET_IFNAME="enp48s3u1u1"
|
||||
export TP_SOCKET_IFNAME="enp48s3u1u1"
|
||||
export HCCL_SOCKET_IFNAME="enp48s3u1u1"
|
||||
export DISAGGREGATED_PREFILL_RANK_TABLE_PATH=path-to-rank-table
|
||||
|
||||
export OMP_PROC_BIND=false
|
||||
export OMP_NUM_THREADS=100
|
||||
|
||||
export VLLM_USE_V1=1
|
||||
|
||||
vllm serve model_path \
|
||||
--host 0.0.0.0 \
|
||||
--port 20002 \
|
||||
--tensor-parallel-size 1\
|
||||
--seed 1024 \
|
||||
--served-model-name dsv3 \
|
||||
--max-model-len 2000 \
|
||||
---max-num-batched-tokens 2000 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector": "LLMDataDistCMgrConnector",
|
||||
"kv_buffer_device": "npu",
|
||||
"kv_role": "kv_consumer",
|
||||
"kv_parallel_size": 1,
|
||||
"kv_port": "20001",
|
||||
"engine_id": 0,
|
||||
"kv_connector_module_path": "vllm_ascend.distributed.llmdatadist_connector_v1_a3"
|
||||
}' \
|
||||
--additional-config \
|
||||
'{"enable_graph_mode": "True"}'\
|
||||
205
examples/eplb/eplb_deepseek.py
Normal file
205
examples/eplb/eplb_deepseek.py
Normal file
@@ -0,0 +1,205 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Expert parallelism load balancer (EPLB) for vLLM.
|
||||
The rearrangement algorithm is adapted from
|
||||
[DeepSeek EPLB](https://github.com/deepseek-ai/eplb).
|
||||
"""
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def balanced_packing(weight: torch.Tensor,
|
||||
num_packs: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Pack n weighted objects to m packs, such that each bin contains exactly n/m objects and the weights of all packs
|
||||
are as balanced as possible.
|
||||
|
||||
Parameters:
|
||||
weight: [X, n], the weight of each item
|
||||
num_packs: number of packs
|
||||
|
||||
Returns:
|
||||
pack_index: [X, n], the pack index of each item
|
||||
rank_in_pack: [X, n], the rank of the item in the pack
|
||||
"""
|
||||
num_layers, num_groups = weight.shape
|
||||
assert num_groups % num_packs == 0
|
||||
groups_per_pack = num_groups // num_packs
|
||||
|
||||
if groups_per_pack == 1:
|
||||
pack_index = torch.arange(weight.size(-1),
|
||||
dtype=torch.int64,
|
||||
device=weight.device).expand(weight.shape)
|
||||
rank_in_pack = torch.zeros_like(weight, dtype=torch.int64)
|
||||
return pack_index, rank_in_pack
|
||||
|
||||
indices = weight.float().sort(-1, descending=True).indices.cpu()
|
||||
pack_index = torch.full_like(weight,
|
||||
fill_value=-1,
|
||||
dtype=torch.int64,
|
||||
device='cpu')
|
||||
rank_in_pack = torch.full_like(pack_index, fill_value=-1)
|
||||
for i in range(num_layers):
|
||||
pack_weights = [0] * num_packs
|
||||
pack_items = [0] * num_packs
|
||||
for group in indices[i]:
|
||||
pack = min(
|
||||
(i
|
||||
for i in range(num_packs) if pack_items[i] < groups_per_pack),
|
||||
key=pack_weights.__getitem__)
|
||||
assert pack_items[pack] < groups_per_pack
|
||||
pack_index[i, group] = pack
|
||||
rank_in_pack[i, group] = pack_items[pack]
|
||||
pack_weights[pack] += weight[i, group]
|
||||
pack_items[pack] += 1
|
||||
return pack_index, rank_in_pack
|
||||
|
||||
|
||||
def replicate_experts(
|
||||
weight: torch.Tensor,
|
||||
num_phy: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Replicate `num_log` experts to `num_phy` replicas, such that the maximum load of all replicas is minimized.
|
||||
|
||||
Parameters:
|
||||
weight: [X, num_log]
|
||||
num_phy: total number of experts after replication
|
||||
|
||||
Returns:
|
||||
phy2log: [X, num_phy], logical expert id of each physical expert
|
||||
rank: [X, num_phy], the replica rank
|
||||
logcnt: [X, num_log], number of replicas for each logical expert
|
||||
"""
|
||||
n, num_log = weight.shape
|
||||
num_redundant = num_phy - num_log
|
||||
assert num_redundant >= 0
|
||||
device = weight.device
|
||||
phy2log = torch.arange(num_phy, dtype=torch.int64,
|
||||
device=device).repeat(n, 1)
|
||||
rank = torch.zeros(n, num_phy, dtype=torch.int64, device=device)
|
||||
logcnt = torch.ones(n, num_log, dtype=torch.int64, device=device)
|
||||
arangen = torch.arange(n, dtype=torch.int64, device=device)
|
||||
for i in range(num_log, num_phy):
|
||||
redundant_indices = (weight / logcnt).max(dim=-1).indices
|
||||
phy2log[:, i] = redundant_indices
|
||||
rank[:, i] = logcnt[arangen, redundant_indices]
|
||||
logcnt[arangen, redundant_indices] += 1
|
||||
return phy2log, rank, logcnt
|
||||
|
||||
|
||||
def rebalance_experts_hierarchical(weight: torch.Tensor,
|
||||
num_physical_experts: int, num_groups: int,
|
||||
num_nodes: int, num_gpus: int):
|
||||
"""
|
||||
Parameters:
|
||||
weight: [num_moe_layers, num_logical_experts]
|
||||
num_physical_experts: number of physical experts after replication
|
||||
num_groups: number of expert groups
|
||||
num_nodes: number of server nodes, where the intra-node network (e.g, NVLink) is faster
|
||||
num_gpus: number of GPUs, must be a multiple of `num_nodes`
|
||||
|
||||
Returns:
|
||||
physical_to_logical_map: [num_moe_layers, num_physical_experts]
|
||||
logical_to_physical_map: [num_moe_layers, num_logical_experts, X]
|
||||
logical_count: [num_moe_layers, num_logical_experts]
|
||||
"""
|
||||
num_layers, num_logical_experts = weight.shape
|
||||
assert num_logical_experts % num_groups == 0
|
||||
group_size = num_logical_experts // num_groups
|
||||
assert num_groups % num_nodes == 0
|
||||
groups_per_node = num_groups // num_nodes
|
||||
assert num_gpus % num_nodes == 0
|
||||
assert num_physical_experts % num_gpus == 0
|
||||
phy_experts_per_gpu = num_physical_experts // num_gpus
|
||||
|
||||
def inverse(perm: torch.Tensor) -> torch.Tensor:
|
||||
inv = torch.empty_like(perm)
|
||||
inv.scatter_(
|
||||
1, perm,
|
||||
torch.arange(perm.size(1), dtype=torch.int64,
|
||||
device=perm.device).expand(perm.shape))
|
||||
return inv
|
||||
|
||||
# Step 1: pack groups to nodes
|
||||
tokens_per_group = weight.unflatten(-1, (num_groups, group_size)).sum(-1)
|
||||
group_pack_index, group_rank_in_pack = balanced_packing(
|
||||
tokens_per_group, num_nodes)
|
||||
log2mlog = (((group_pack_index * groups_per_node + group_rank_in_pack) *
|
||||
group_size).unsqueeze(-1) +
|
||||
torch.arange(group_size,
|
||||
dtype=torch.int64,
|
||||
device=group_pack_index.device)).flatten(-2)
|
||||
mlog2log = inverse(log2mlog)
|
||||
|
||||
# Step 2: construct redundant experts within nodes
|
||||
# [num_layers * num_nodes, num_logical_experts // num_nodes]
|
||||
tokens_per_mlog = weight.gather(-1, mlog2log).view(
|
||||
-1, num_logical_experts // num_nodes)
|
||||
phy2mlog, phyrank, mlogcnt = replicate_experts(
|
||||
tokens_per_mlog, num_physical_experts // num_nodes)
|
||||
|
||||
# Step 3: pack physical_experts to GPUs
|
||||
# [num_layers * num_nodes, num_physical_experts // num_nodes]
|
||||
tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog)
|
||||
pack_index, rank_in_pack = balanced_packing(tokens_per_phy,
|
||||
num_gpus // num_nodes)
|
||||
phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack
|
||||
pphy2phy = inverse(phy2pphy)
|
||||
|
||||
pphy2mlog = phy2mlog.gather(
|
||||
-1, pphy2phy) # [num_layers * num_nodes, num_log_per_nodes]
|
||||
pphy2mlog = (pphy2mlog.view(num_layers, num_nodes, -1) + torch.arange(
|
||||
0,
|
||||
num_logical_experts,
|
||||
num_logical_experts // num_nodes,
|
||||
device=group_pack_index.device).view(1, -1, 1)).flatten(-2)
|
||||
pphy2log = mlog2log.gather(-1, pphy2mlog)
|
||||
pphyrank = phyrank.gather(-1, pphy2phy).view(num_layers, -1)
|
||||
logcnt = mlogcnt.view(num_layers, -1).gather(-1, log2mlog)
|
||||
return pphy2log, pphyrank, logcnt
|
||||
|
||||
|
||||
def rebalance_experts(
|
||||
weight: torch.Tensor, num_replicas: int, num_groups: int,
|
||||
num_nodes: int,
|
||||
num_gpus: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Entry point for expert-parallelism load balancer.
|
||||
|
||||
Parameters:
|
||||
weight: [layers, num_logical_experts], the load statistics for all logical experts
|
||||
num_replicas: number of physical experts, must be a multiple of `num_gpus`
|
||||
num_groups: number of expert groups
|
||||
num_nodes: number of server nodes, where the intra-node network (e.g, NVLink) is faster
|
||||
num_gpus: number of GPUs, must be a multiple of `num_nodes`
|
||||
|
||||
Returns:
|
||||
physical_to_logical_map: [layers, num_replicas], the expert index of each replica
|
||||
logical_to_physical_map: [layers, num_logical_experts, X], the replica indices for each expert
|
||||
expert_count: [layers, num_logical_experts], number of physical replicas for each logical expert
|
||||
"""
|
||||
num_layers, num_logical_experts = weight.shape
|
||||
weight = weight.float().cpu()
|
||||
if num_groups % num_nodes == 0:
|
||||
# use hierarchical load-balance policy
|
||||
phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
|
||||
weight, num_replicas, num_groups, num_nodes, num_gpus)
|
||||
else:
|
||||
# use global load-balance policy
|
||||
phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
|
||||
weight, num_replicas, 1, 1, num_gpus)
|
||||
maxlogcnt = logcnt.max().item()
|
||||
log2phy: torch.Tensor = torch.full(
|
||||
(num_layers, num_logical_experts, maxlogcnt),
|
||||
-1,
|
||||
dtype=torch.int64,
|
||||
device=logcnt.device)
|
||||
log2phy.view(num_layers, -1).scatter_(
|
||||
-1, phy2log * maxlogcnt + phyrank,
|
||||
torch.arange(num_replicas, dtype=torch.int64,
|
||||
device=log2phy.device).expand(num_layers, -1))
|
||||
return phy2log, log2phy, logcnt
|
||||
|
||||
|
||||
__all__ = ['rebalance_experts']
|
||||
186
examples/eplb/eplb_strategy.py
Normal file
186
examples/eplb/eplb_strategy.py
Normal file
@@ -0,0 +1,186 @@
|
||||
# coding=utf-8
|
||||
# Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
|
||||
import matplotlib.pyplot as plt # type: ignore
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
logger = logging.getLogger("msit_logger")
|
||||
|
||||
|
||||
def save_matrix_to_json(output_path, file_name, deployment):
|
||||
num_layers = deployment.shape[0]
|
||||
num_cards = deployment.shape[1]
|
||||
|
||||
data = {"moe_layer_count": num_layers}
|
||||
layer_list = []
|
||||
for i in range(num_layers):
|
||||
layer = {"layer_id": i, "device_count": num_cards}
|
||||
device_list = []
|
||||
for j in range(num_cards):
|
||||
device = {
|
||||
"device_id": j,
|
||||
"device_expert": deployment[i, j].tolist()
|
||||
}
|
||||
device_list.append(device)
|
||||
layer["device_list"] = device_list
|
||||
layer_list.append(layer)
|
||||
data["layer_list"] = layer_list
|
||||
|
||||
file_name = f"{output_path}{file_name}.json"
|
||||
|
||||
# Save as JSON file
|
||||
try:
|
||||
with open(file_name, 'w') as f:
|
||||
json.dump(data, f, indent=4)
|
||||
except Exception as e:
|
||||
print(f"write {file_name} failed: {e}")
|
||||
|
||||
|
||||
def calculate_average(lst):
|
||||
"""calculate the average of a list"""
|
||||
if not lst:
|
||||
raise ValueError("list is empty")
|
||||
|
||||
total = 0.0
|
||||
count = 0
|
||||
|
||||
for element in lst:
|
||||
# Check if element is numeric
|
||||
if isinstance(element, (int, float, np.int64, np.float64)):
|
||||
total += float(element)
|
||||
count += 1
|
||||
else:
|
||||
# Non-numeric elements will be ignored with a warning
|
||||
print(f"warning: element {element} is not a number, ignored")
|
||||
|
||||
if count == 0:
|
||||
raise ValueError("list does not contain any number")
|
||||
|
||||
return total / count
|
||||
|
||||
|
||||
def layer_imblance_polt(y_list, label_names, device_num, output_path,
|
||||
file_name):
|
||||
|
||||
plt.rcParams['font.sans-serif'] = ['Arial']
|
||||
plt.rcParams['axes.unicode_minus'] = False
|
||||
x = [i for i in range(58)]
|
||||
for index, y in enumerate(y_list):
|
||||
plt.plot(x,
|
||||
y,
|
||||
label=rf'{label_names[index]},avg={calculate_average(y)}')
|
||||
|
||||
plt.legend()
|
||||
plt.title(rf'Load Distribution (num_gpus={device_num})')
|
||||
plt.xlabel('layer')
|
||||
plt.ylabel('Device Load')
|
||||
|
||||
# Show grid lines
|
||||
plt.grid(True)
|
||||
|
||||
plt.savefig(os.path.join(output_path, file_name), dpi=300)
|
||||
|
||||
# Clear current plot
|
||||
plt.close()
|
||||
|
||||
|
||||
def deepseek_deploy(workload, num_redundancy_expert, num_groups, num_nodes,
|
||||
num_gpus, num_original_expert):
|
||||
from eplb_deepseek import rebalance_experts
|
||||
num_replicas = num_original_expert + num_redundancy_expert
|
||||
hy2log, log2phy, logcnt = rebalance_experts(workload, num_replicas,
|
||||
num_groups, num_nodes,
|
||||
num_gpus)
|
||||
|
||||
# Convert to global_deployment
|
||||
workload = workload.cpu().numpy()
|
||||
global_deployment = []
|
||||
layer_num = log2phy.shape[0]
|
||||
num_physical_experts_local = (num_original_expert +
|
||||
num_redundancy_expert) // num_gpus
|
||||
for layer_idx in range(layer_num):
|
||||
layer_deployment = []
|
||||
for gpu_idx in range(num_gpus):
|
||||
local_deployment = hy2log[layer_idx][gpu_idx *
|
||||
num_physical_experts_local:
|
||||
(gpu_idx + 1) *
|
||||
num_physical_experts_local]
|
||||
local_deployment = local_deployment.flatten()
|
||||
layer_deployment.append(local_deployment.tolist())
|
||||
global_deployment.append(layer_deployment)
|
||||
|
||||
# Remap expert distribution according to log2phy
|
||||
original_weights = []
|
||||
max_weights = []
|
||||
average_weights = []
|
||||
y_list = []
|
||||
for layer_idx in range(layer_num):
|
||||
new_value = workload[layer_idx].reshape(num_gpus, -1)
|
||||
row_sum = np.sum(new_value, axis=1)
|
||||
original_weights.append(row_sum.max())
|
||||
average_weights.append((np.sum(workload[layer_idx]) / num_gpus))
|
||||
|
||||
opt_workload = np.zeros((num_original_expert + num_redundancy_expert),
|
||||
dtype=np.float64)
|
||||
for expert_idx in range(num_original_expert):
|
||||
physical_expert_idxs = log2phy[layer_idx][expert_idx]
|
||||
physical_expert_idxs = physical_expert_idxs.flatten()
|
||||
physical_expert_idxs = physical_expert_idxs[
|
||||
physical_expert_idxs != -1]
|
||||
for physical_expert_idx in physical_expert_idxs:
|
||||
opt_workload[physical_expert_idx] += workload[layer_idx][
|
||||
expert_idx] / len(physical_expert_idxs)
|
||||
opt_workload = opt_workload.reshape(num_gpus, -1)
|
||||
row_sum = np.sum(opt_workload, axis=1)
|
||||
max_weights.append(row_sum.max())
|
||||
|
||||
y_list = [original_weights, max_weights, average_weights]
|
||||
return global_deployment, y_list
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--exp_name", type=str, default="gsm8k_temp0.0")
|
||||
parser.add_argument("--num_original_expert", type=int, default=256)
|
||||
parser.add_argument("--input_path", type=str, default="")
|
||||
parser.add_argument("--output_path", type=str, default="")
|
||||
parser.add_argument("--num_redundancy_expert", type=int, default=0)
|
||||
parser.add_argument("--num_devices", type=int, default=32)
|
||||
parser.add_argument("--num_groups", type=int, default=8)
|
||||
parser.add_argument("--num_nodes", type=int, default=4)
|
||||
args = parser.parse_args()
|
||||
exp_name = args.exp_name
|
||||
input_path = args.input_path
|
||||
output_path = args.output_path
|
||||
os.makedirs(output_path, exist_ok=True)
|
||||
num_redundancy_expert = args.num_redundancy_expert
|
||||
num_devices = args.num_devices
|
||||
num_original_expert = args.num_original_expert
|
||||
num_groups = args.num_groups
|
||||
num_nodes = args.num_nodes
|
||||
|
||||
# NOTE: assume input workload format: [layer_num, num_experts]
|
||||
workload = torch.load(input_path, map_location=torch.device('cpu'))
|
||||
global_deployment, y_list = deepseek_deploy(workload,
|
||||
num_redundancy_expert,
|
||||
num_groups, num_nodes,
|
||||
num_devices,
|
||||
num_original_expert)
|
||||
|
||||
file_name = f"{exp_name}_{num_devices}_{num_redundancy_expert}"
|
||||
save_matrix_to_json(output_path, file_name, np.array(global_deployment))
|
||||
label_names = [
|
||||
'default deployment max load', 'balanced load max load',
|
||||
'balanced load avg load'
|
||||
]
|
||||
new_file_name = f"{exp_name}_{num_devices}_{num_redundancy_expert}.png"
|
||||
layer_imblance_polt(y_list, label_names, num_devices, output_path,
|
||||
new_file_name)
|
||||
38
examples/external_online_dp/README.md
Normal file
38
examples/external_online_dp/README.md
Normal file
@@ -0,0 +1,38 @@
|
||||
Here is an example guiding how to use `launch_online_dp.py` to launch external dp server in vllm. User can easily launch external dp server following the steps below:
|
||||
|
||||
### Modify parameters in `run_dp_template.sh`
|
||||
`run_dp_template.sh` is an template script used to launch each dp vllm instance separately. It will be called by `launch_online_dp.py` in multi threads and most of its configurations are set by `launch_online_dp.py`. Parameters you need to set manually include:
|
||||
|
||||
1. The IP and socket_ifname of your machine. If running on multi-nodes, please make sure the scripts on each node has been set with correct IP and socket_ifname of that node.
|
||||
2. vLLM serving related parameters including model_path and other configurations. Note that port, dp-related parammeters and tp_size is set by `launch_online_dp.py`, all the other vLLM parameters in this file only serve as an example and you are free to modify them according to your purpose.
|
||||
|
||||
### Run `launch_online_dp.py` with CL arguments
|
||||
All the arguments that can be set by users are:
|
||||
|
||||
1. `--dp-size`: global data parallel size, must be set
|
||||
2. `--tp-size`: tensor parallel size, default 1
|
||||
3. `--dp-size-local`: local data parallel size, defaultly set to `dp_size`
|
||||
4. `--dp-rank-start`: Starting rank for data parallel, default 0
|
||||
5. `--dp-address`: IP address of data parallel master node
|
||||
6. `--dp-rpc-port`: Port of data parallel master node, default 12345
|
||||
7. `--vllm-start-port`: Starting port of vLLM serving instances, default 9000
|
||||
|
||||
An example of running external DP in one single node:
|
||||
```(python)
|
||||
cd examples/external_online_dp
|
||||
# running DP4 TP4 in a node with 16 NPUs
|
||||
python launch_online_dp.py --dp-size 4 --tp-size 4 --dp-size-local 4 --dp-rank-start 0 --dp-address x.x.x.x --dp-rpc-port 12342
|
||||
```
|
||||
|
||||
An example of running external DP in two nodes:
|
||||
```(python)
|
||||
cd examples/external_online_dp
|
||||
# running DP4 TP4 in two nodes with 8 NPUs each
|
||||
|
||||
# On node 0:
|
||||
python launch_online_dp.py --dp-size 4 --tp-size 4 --dp-size-local 2 --dp-rank-start 0 --dp-address x.x.x.x --dp-rpc-port 12342
|
||||
|
||||
# On node 1:
|
||||
python launch_online_dp.py --dp-size 4 --tp-size 4 --dp-size-local 2 --dp-rank-start 2 --dp-address x.x.x.x --dp-rpc-port 12342
|
||||
```
|
||||
|
||||
97
examples/external_online_dp/launch_online_dp.py
Normal file
97
examples/external_online_dp/launch_online_dp.py
Normal file
@@ -0,0 +1,97 @@
|
||||
import argparse
|
||||
import multiprocessing
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--dp-size",
|
||||
type=int,
|
||||
required=True,
|
||||
help="Data parallel size."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tp-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Tensor parallel size."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dp-size-local",
|
||||
type=int,
|
||||
default=-1,
|
||||
help="Local data parallel size."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dp-rank-start",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Starting rank for data parallel."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dp-address",
|
||||
type=str,
|
||||
required=True,
|
||||
help="IP address for data parallel master node."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dp-rpc-port",
|
||||
type=str,
|
||||
default=12345,
|
||||
help="Port for data parallel master node."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vllm-start-port",
|
||||
type=int,
|
||||
default=9000,
|
||||
help="Starting port for the engine."
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
args = parse_args()
|
||||
dp_size = args.dp_size
|
||||
tp_size = args.tp_size
|
||||
dp_size_local = args.dp_size_local
|
||||
if dp_size_local == -1:
|
||||
dp_size_local = dp_size
|
||||
dp_rank_start = args.dp_rank_start
|
||||
dp_address = args.dp_address
|
||||
dp_rpc_port = args.dp_rpc_port
|
||||
vllm_start_port = args.vllm_start_port
|
||||
|
||||
def run_command(visiable_devices, dp_rank, vllm_engine_port):
|
||||
command = [
|
||||
"bash",
|
||||
"./run_dp_template.sh",
|
||||
visiable_devices,
|
||||
str(vllm_engine_port),
|
||||
str(dp_size),
|
||||
str(dp_rank),
|
||||
dp_address,
|
||||
dp_rpc_port,
|
||||
str(tp_size),
|
||||
]
|
||||
subprocess.run(command, check=True)
|
||||
|
||||
if __name__ == "__main__":
|
||||
template_path = "./run_dp_template.sh"
|
||||
if not os.path.exists(template_path):
|
||||
print(f"Template file {template_path} does not exist.")
|
||||
sys.exit(1)
|
||||
|
||||
processes = []
|
||||
num_cards = dp_size_local * tp_size
|
||||
for i in range(dp_size_local):
|
||||
dp_rank = dp_rank_start + i
|
||||
vllm_engine_port = vllm_start_port + i
|
||||
visiable_devices = ",".join(str(x) for x in range(i * tp_size, (i + 1) * tp_size))
|
||||
process = multiprocessing.Process(target=run_command,
|
||||
args=(visiable_devices, dp_rank,
|
||||
vllm_engine_port))
|
||||
processes.append(process)
|
||||
process.start()
|
||||
|
||||
for process in processes:
|
||||
process.join()
|
||||
46
examples/external_online_dp/run_dp_template.sh
Normal file
46
examples/external_online_dp/run_dp_template.sh
Normal file
@@ -0,0 +1,46 @@
|
||||
export HCCL_IF_IP=your_ip_here
|
||||
export GLOO_SOCKET_IFNAME=your_socket_ifname_here
|
||||
export TP_SOCKET_IFNAME=your_socket_ifname_here
|
||||
export HCCL_SOCKET_IFNAME=your_socket_ifname_here
|
||||
export DISAGGREGATED_PREFILL_RANK_TABLE_PATH=your_rank_table_path_here
|
||||
export VLLM_LOGGING_LEVEL="info"
|
||||
export OMP_PROC_BIND=false
|
||||
export OMP_NUM_THREADS=10
|
||||
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
|
||||
export HCCL_DETERMINISTIC=True
|
||||
export HCCL_BUFFSIZE=1024
|
||||
export TASK_QUEUE_ENABLE=1
|
||||
|
||||
export VLLM_USE_V1=1
|
||||
|
||||
export ASCEND_RT_VISIBLE_DEVICES=$1
|
||||
|
||||
vllm serve model_path \
|
||||
--host 0.0.0.0 \
|
||||
--port $2 \
|
||||
--data-parallel-size $3 \
|
||||
--data-parallel-rank $4 \
|
||||
--data-parallel-address $5 \
|
||||
--data-parallel-rpc-port $6 \
|
||||
--tensor-parallel-size $7 \
|
||||
--enable-expert-parallel \
|
||||
--seed 1024 \
|
||||
--served-model-name dsv3 \
|
||||
--max-model-len 3500 \
|
||||
--max-num-batched-tokens 3500 \
|
||||
--max-num-seqs 28 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--quantization ascend \
|
||||
--speculative-config '{"num_speculative_tokens": 1, "method":"deepseek_mtp"}' \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector": "LLMDataDistCMgrConnector",
|
||||
"kv_buffer_device": "npu",
|
||||
"kv_role": "kv_consumer",
|
||||
"kv_parallel_size": "1",
|
||||
"kv_port": "20001",
|
||||
"engine_id": "0",
|
||||
"kv_connector_module_path": "vllm_ascend.distributed.llmdatadist_c_mgr_connector"
|
||||
}' \
|
||||
--additional-config \
|
||||
'{"ascend_scheduler_config": {"enabled": true}, "torchair_graph_config":{"enabled":true,"enable_kv_nz":false, "enable_multistream_moe":false, "graph_batch_size":[28]}, "enable_weight_nz_layout":true}'
|
||||
257
examples/offline_data_parallel.py
Normal file
257
examples/offline_data_parallel.py
Normal file
@@ -0,0 +1,257 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
# Adapted from vllm-project/vllm/examples/offline_inference/data_parallel.py
|
||||
#
|
||||
"""
|
||||
Usage:
|
||||
Single node:
|
||||
Dense models:
|
||||
python examples/offline_data_parallel.py \
|
||||
--model="Qwen/Qwen2.5-0.5B-Instruct" \
|
||||
--dp-size=2 \
|
||||
--tp-size=2
|
||||
MOE models:
|
||||
python examples/offline_data_parallel.py \
|
||||
--model="ibm-research/PowerMoE-3b" \
|
||||
--dp-size=2 \
|
||||
--tp-size=2 \
|
||||
--enable-expert-parallel
|
||||
|
||||
Multi-node:
|
||||
Node 0 (assume the node has ip of 10.99.48.128):
|
||||
python examples/offline_data_parallel.py \
|
||||
--model="ibm-research/PowerMoE-3b" \
|
||||
--dp-size=2 \
|
||||
--tp-size=2 \
|
||||
--node-size=2 \
|
||||
--node-rank=0 \
|
||||
--enable-expert-parallel \
|
||||
--master-addr=10.99.48.128 \
|
||||
--master-port=13345
|
||||
Node 1:
|
||||
python examples/offline_data_parallel.py \
|
||||
--model="ibm-research/PowerMoE-3b" \
|
||||
--dp-size=2 \
|
||||
--tp-size=2 \
|
||||
--node-size=2 \
|
||||
--node-rank=1 \
|
||||
--enable-expert-parallel \
|
||||
--master-addr=10.99.48.128 \
|
||||
--master-port=13345
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import gc
|
||||
import os
|
||||
from time import sleep
|
||||
|
||||
import torch
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.distributed.parallel_state import ( # noqa E402
|
||||
destroy_distributed_environment, destroy_model_parallel)
|
||||
from vllm.utils import get_open_port
|
||||
|
||||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
def parse_args():
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Data Parallel Inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="ibm-research/PowerMoE-3b",
|
||||
help="Model name or path",
|
||||
)
|
||||
parser.add_argument("--dp-size",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Data parallel size")
|
||||
parser.add_argument("--tp-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Tensor parallel size")
|
||||
parser.add_argument("--node-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Total number of nodes")
|
||||
parser.add_argument("--node-rank",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Rank of the current node")
|
||||
parser.add_argument("--master-addr",
|
||||
type=str,
|
||||
default="",
|
||||
help="Master node IP address")
|
||||
parser.add_argument("--master-port",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Master node port")
|
||||
parser.add_argument("--enforce-eager",
|
||||
action="store_true",
|
||||
help="Enforce eager mode execution.")
|
||||
parser.add_argument("--trust-remote-code",
|
||||
action="store_true",
|
||||
help="Trust remote code.")
|
||||
parser.add_argument("--enable-expert-parallel",
|
||||
action="store_true",
|
||||
help="Enable expert parallel, used in MOE models.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def cleanup_env_and_memory():
|
||||
destroy_model_parallel()
|
||||
destroy_distributed_environment()
|
||||
with contextlib.suppress(AssertionError):
|
||||
torch.distributed.destroy_process_group()
|
||||
gc.collect()
|
||||
torch.npu.empty_cache()
|
||||
torch.npu.reset_peak_memory_stats()
|
||||
|
||||
def main(
|
||||
model,
|
||||
dp_size,
|
||||
local_dp_rank,
|
||||
global_dp_rank,
|
||||
dp_master_ip,
|
||||
dp_master_port,
|
||||
GPUs_per_dp_rank,
|
||||
enable_expert_parallel,
|
||||
enforce_eager,
|
||||
trust_remote_code,
|
||||
):
|
||||
# DP only support on V1 engine
|
||||
os.environ["VLLM_DP_RANK"] = str(global_dp_rank)
|
||||
os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank)
|
||||
os.environ["VLLM_DP_SIZE"] = str(dp_size)
|
||||
os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip
|
||||
os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)
|
||||
|
||||
# CUDA_VISIBLE_DEVICES for each DP rank is set automatically inside the
|
||||
# engine processes.
|
||||
|
||||
# Sample prompts.
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
] * 100
|
||||
|
||||
# with DP, each rank should process different prompts.
|
||||
# usually all the DP ranks process a full dataset,
|
||||
# and each rank processes a different part of the dataset.
|
||||
floor = len(prompts) // dp_size
|
||||
remainder = len(prompts) % dp_size
|
||||
|
||||
# Distribute prompts into even groups.
|
||||
def start(rank):
|
||||
return rank * floor + min(rank, remainder)
|
||||
|
||||
prompts = prompts[start(global_dp_rank):start(global_dp_rank + 1)]
|
||||
if len(prompts) == 0:
|
||||
# if any rank has no prompts to process,
|
||||
# we need to set a placeholder prompt
|
||||
prompts = ["Placeholder"]
|
||||
print(f"DP rank {global_dp_rank} needs to process {len(prompts)} prompts")
|
||||
|
||||
# Create a sampling params object.
|
||||
# since we are doing data parallel, every rank can have different
|
||||
# sampling params. here we set different max_tokens for different
|
||||
# ranks for demonstration.
|
||||
sampling_params = SamplingParams(temperature=0.8,
|
||||
top_p=0.95,
|
||||
max_tokens=[16, 20][global_dp_rank % 2])
|
||||
|
||||
# Create an LLM.
|
||||
llm = LLM(
|
||||
model=model,
|
||||
tensor_parallel_size=GPUs_per_dp_rank,
|
||||
enforce_eager=enforce_eager,
|
||||
enable_expert_parallel=enable_expert_parallel,
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
# Print the outputs.
|
||||
for i, output in enumerate(outputs):
|
||||
if i >= 5:
|
||||
# print only 5 outputs
|
||||
break
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"DP rank {global_dp_rank}, Prompt: {prompt!r}, "
|
||||
f"Generated text: {generated_text!r}")
|
||||
|
||||
# Give engines time to pause their processing loops before exiting.
|
||||
sleep(5)
|
||||
del llm
|
||||
cleanup_env_and_memory()
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
|
||||
dp_size = args.dp_size
|
||||
tp_size = args.tp_size
|
||||
node_size = args.node_size
|
||||
node_rank = args.node_rank
|
||||
|
||||
if node_size == 1:
|
||||
dp_master_ip = "127.0.0.1"
|
||||
dp_master_port = get_open_port()
|
||||
else:
|
||||
dp_master_ip = args.master_addr
|
||||
dp_master_port = args.master_port
|
||||
|
||||
assert dp_size % node_size == 0, "dp_size should be divisible by node_size"
|
||||
dp_per_node = dp_size // node_size
|
||||
|
||||
from multiprocessing import Process
|
||||
|
||||
procs = []
|
||||
for local_dp_rank, global_dp_rank in enumerate(
|
||||
range(node_rank * dp_per_node, (node_rank + 1) * dp_per_node)):
|
||||
proc = Process(
|
||||
target=main,
|
||||
args=(
|
||||
args.model,
|
||||
dp_size,
|
||||
local_dp_rank,
|
||||
global_dp_rank,
|
||||
dp_master_ip,
|
||||
dp_master_port,
|
||||
tp_size,
|
||||
args.enable_expert_parallel,
|
||||
args.enforce_eager,
|
||||
args.trust_remote_code,
|
||||
),
|
||||
)
|
||||
proc.start()
|
||||
procs.append(proc)
|
||||
exit_code = 0
|
||||
for proc in procs:
|
||||
proc.join(timeout=300)
|
||||
if proc.exitcode is None:
|
||||
print(
|
||||
f"Killing process {proc.pid} that didn't stop within 5 minutes."
|
||||
)
|
||||
proc.kill()
|
||||
exit_code = 1
|
||||
elif proc.exitcode:
|
||||
exit_code = proc.exitcode
|
||||
|
||||
exit(exit_code)
|
||||
147
examples/offline_disaggregated_prefill_npu.py
Normal file
147
examples/offline_disaggregated_prefill_npu.py
Normal file
@@ -0,0 +1,147 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
# Adapted from vllm-project/vllm/examples/offline_inference/basic.py
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import time
|
||||
from multiprocessing import Event, Process
|
||||
|
||||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
def clean_up():
|
||||
import gc
|
||||
|
||||
import torch
|
||||
from vllm.distributed.parallel_state import (
|
||||
destroy_distributed_environment, destroy_model_parallel)
|
||||
destroy_model_parallel()
|
||||
destroy_distributed_environment()
|
||||
gc.collect()
|
||||
torch.npu.empty_cache()
|
||||
|
||||
|
||||
def run_prefill(prefill_done, process_close):
|
||||
# ranktable.json needs be generated using gen_ranktable.sh
|
||||
# from the examples/disaggregated_prefill_v1 in the main branch.
|
||||
os.environ['DISAGGREGATED_PREFILL_RANK_TABLE_PATH'] = "./ranktable.json"
|
||||
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = "0"
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import KVTransferConfig
|
||||
|
||||
prompts = [
|
||||
"Hello, how are you today?", "Hi, what is your name?",
|
||||
"Tell me a very long story.", "what is your favourite book?"
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
|
||||
|
||||
ktc = KVTransferConfig(kv_connector="LLMDataDistCMgrConnector", kv_buffer_device="npu", kv_role="kv_producer",
|
||||
kv_parallel_size=1,
|
||||
kv_connector_module_path="vllm_ascend.distributed.llmdatadist_c_mgr_connector")
|
||||
# Set NPU memory utilization to 0.8
|
||||
llm = LLM(model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
|
||||
kv_transfer_config=ktc,
|
||||
max_model_len=2000,
|
||||
gpu_memory_utilization=0.8,
|
||||
tensor_parallel_size=1)
|
||||
|
||||
llm.generate(prompts, sampling_params)
|
||||
print("Prefill node is finished.")
|
||||
prefill_done.set()
|
||||
|
||||
# To keep the prefill node running in case the decode node is not done
|
||||
# otherwise, the script might exit prematurely, causing incomplete decoding.
|
||||
try:
|
||||
while not process_close.is_set():
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
print("Script stopped by user.")
|
||||
finally:
|
||||
print("Cleanup prefill resources")
|
||||
del llm
|
||||
clean_up()
|
||||
|
||||
|
||||
def run_decode(prefill_done):
|
||||
os.environ['VLLM_LLMDD_RPC_PORT'] = '6634'
|
||||
# ranktable.json needs be generated using gen_ranktable.sh
|
||||
# from the examples/disaggregated_prefill_v1 module in the main branch.
|
||||
os.environ['DISAGGREGATED_PREFILL_RANK_TABLE_PATH'] = "./ranktable.json"
|
||||
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = "1"
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import KVTransferConfig
|
||||
|
||||
prompts = [
|
||||
"Hello, how are you today?", "Hi, what is your name?",
|
||||
"Tell me a very long story.", "what is your favourite book?"
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0, top_p=0.95)
|
||||
|
||||
ktc = KVTransferConfig(kv_connector="LLMDataDistCMgrConnector", kv_buffer_device="npu", kv_role="kv_consumer",
|
||||
kv_parallel_size=1, kv_connector_module_path="vllm_ascend.distributed.llmdatadist_c_mgr_connector")
|
||||
|
||||
llm = LLM(model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
|
||||
kv_transfer_config=ktc,
|
||||
max_model_len=2000,
|
||||
gpu_memory_utilization=0.8,
|
||||
tensor_parallel_size=1)
|
||||
|
||||
# Wait for the producer to start the consumer
|
||||
print("Waiting for prefill node to finish...")
|
||||
prefill_done.wait()
|
||||
|
||||
# At this point when the prefill_done is set, the kv-cache should have been
|
||||
# transferred to this decode node, so we can start decoding.
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
del llm
|
||||
clean_up()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
mp.get_context('spawn')
|
||||
|
||||
prefill_done = Event()
|
||||
process_close = Event()
|
||||
prefill_process = Process(target=run_prefill,
|
||||
args=(
|
||||
prefill_done,
|
||||
process_close,
|
||||
))
|
||||
decode_process = Process(target=run_decode, args=(prefill_done, ))
|
||||
|
||||
# Start prefill node
|
||||
prefill_process.start()
|
||||
|
||||
# Start decode node
|
||||
decode_process.start()
|
||||
|
||||
# Terminate the prefill node when decode is finished
|
||||
decode_process.join()
|
||||
|
||||
# Terminate prefill process
|
||||
process_close.set()
|
||||
prefill_process.join()
|
||||
prefill_process.terminate()
|
||||
print("All process done!")
|
||||
52
examples/offline_dualbatch_overlap_npu.py
Normal file
52
examples/offline_dualbatch_overlap_npu.py
Normal file
@@ -0,0 +1,52 @@
|
||||
import os
|
||||
import time
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
# enable dual-batch overlap for vllm ascend
|
||||
os.environ["VLLM_ASCEND_ENABLE_DBO"] = "1"
|
||||
|
||||
# Sample prompts.
|
||||
prompts = ["The president of the United States is"] * 41
|
||||
# Create a sampling params object.
|
||||
sampling_params = SamplingParams(max_tokens=100, temperature=0.0)
|
||||
|
||||
|
||||
def main():
|
||||
# Create an LLM.
|
||||
llm = LLM(model="deepseek-ai/DeepSeek-V3-Lite-base-latest-w8a8-dynamic",
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=2,
|
||||
max_model_len=4096,
|
||||
trust_remote_code=True,
|
||||
enable_expert_parallel=True,
|
||||
additional_config={
|
||||
"torchair_graph_config": {
|
||||
"enabled": False
|
||||
},
|
||||
"ascend_scheduler_config": {
|
||||
"enabled": True
|
||||
},
|
||||
})
|
||||
|
||||
# Generate texts from the prompts. The output is a list of RequestOutput
|
||||
# objects that contain the prompt, generated text, and other information.
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
# Print the outputs.
|
||||
print("-" * 50)
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
|
||||
print("-" * 50)
|
||||
|
||||
# Add a buffer to wait for profiler in the background process
|
||||
# (in case MP is on) to finish writing profiling output.
|
||||
time.sleep(10)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
58
examples/offline_embed.py
Normal file
58
examples/offline_embed.py
Normal file
@@ -0,0 +1,58 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
# Adapted from https://www.modelscope.cn/models/Qwen/Qwen3-Embedding-0.6B
|
||||
#
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
from vllm import LLM
|
||||
|
||||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
def get_detailed_instruct(task_description: str, query: str) -> str:
|
||||
return f'Instruct: {task_description}\nQuery:{query}'
|
||||
|
||||
|
||||
def main():
|
||||
# Each query must come with a one-sentence instruction that describes the task
|
||||
task = 'Given a web search query, retrieve relevant passages that answer the query'
|
||||
|
||||
queries = [
|
||||
get_detailed_instruct(task, 'What is the capital of China?'),
|
||||
get_detailed_instruct(task, 'Explain gravity')
|
||||
]
|
||||
# No need to add instruction for retrieval documents
|
||||
documents = [
|
||||
"The capital of China is Beijing.",
|
||||
"Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun."
|
||||
]
|
||||
input_texts = queries + documents
|
||||
|
||||
model = LLM(model="Qwen/Qwen3-Embedding-0.6B", task="embed")
|
||||
|
||||
outputs = model.embed(input_texts)
|
||||
embeddings = torch.tensor([o.outputs.embedding for o in outputs])
|
||||
# Calculate the similarity scores between the first two queries and the last two documents
|
||||
scores = (embeddings[:2] @ embeddings[2:].T)
|
||||
print(scores.tolist())
|
||||
# [[0.7620252966880798, 0.14078938961029053], [0.1358368694782257, 0.6013815999031067]]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
287
examples/offline_external_launcher.py
Normal file
287
examples/offline_external_launcher.py
Normal file
@@ -0,0 +1,287 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
# Adapted from vllm-project/vllm/examples/offline_inference/data_parallel.py
|
||||
|
||||
# Note: This script is designed to run with e2e test,
|
||||
# please be careful to modify it.
|
||||
"""
|
||||
Usage:
|
||||
Single node:
|
||||
Dense models:
|
||||
python examples/offline_external_launcher.py \
|
||||
--model="Qwen/Qwen2.5-0.5B-Instruct" \
|
||||
--tp-size=1 \
|
||||
--proc-per-node=2
|
||||
MOE models:
|
||||
python examples/offline_external_launcher.py \
|
||||
--model="Qwen/Qwen3-30B-A3B" \
|
||||
--tp-size=2 \
|
||||
--proc-per-node=2 \
|
||||
--enable-expert-parallel
|
||||
|
||||
Multi-node:
|
||||
Node 0 (assume the node has ip of 10.99.48.128):
|
||||
python examples/offline_external_launcher.py \
|
||||
--model="Qwen/Qwen3-30B-A3B" \
|
||||
--tp-size=2 \
|
||||
--node-size=2 \
|
||||
--node-rank=0 \
|
||||
--proc-per-node=2 \
|
||||
--enable-expert-parallel \
|
||||
--master-addr=10.99.48.128 \
|
||||
--master-port=13345
|
||||
Node 1:
|
||||
python examples/offline_external_launcher.py \
|
||||
--model="Qwen/Qwen3-30B-A3B" \
|
||||
--tp-size=2 \
|
||||
--node-size=2 \
|
||||
--node-rank=1 \
|
||||
--enable-expert-parallel \
|
||||
--master-addr=10.99.48.128 \
|
||||
--master-port=13345
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import gc
|
||||
import os
|
||||
from multiprocessing import Process
|
||||
from time import sleep
|
||||
|
||||
import torch
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.distributed.parallel_state import ( # noqa E402
|
||||
destroy_distributed_environment, destroy_model_parallel, get_tp_group)
|
||||
from vllm.utils import get_open_port, GiB_bytes
|
||||
|
||||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
|
||||
def parse_args():
|
||||
|
||||
parser = argparse.ArgumentParser(description="External launcher Inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="Qwen/Qwen3-0.6B",
|
||||
help="Model name or path",
|
||||
)
|
||||
parser.add_argument("--tp-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Tensor parallel size")
|
||||
parser.add_argument("--node-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Total number of nodes")
|
||||
parser.add_argument("--node-rank",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Rank of the current node")
|
||||
parser.add_argument("--proc-per-node",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of processes per node")
|
||||
parser.add_argument("--master-addr",
|
||||
type=str,
|
||||
default="",
|
||||
help="Master node IP address")
|
||||
parser.add_argument("--master-port",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Master node port")
|
||||
parser.add_argument("--enforce-eager",
|
||||
action="store_true",
|
||||
help="Enforce eager mode execution.")
|
||||
parser.add_argument("--trust-remote-code",
|
||||
action="store_true",
|
||||
help="Trust remote code.")
|
||||
parser.add_argument("--enable-expert-parallel",
|
||||
action="store_true",
|
||||
help="Enable expert parallel, used in MOE models.")
|
||||
parser.add_argument("--enable-sleep-mode",
|
||||
action="store_true",
|
||||
help="Enable sleep mode for the engine.")
|
||||
parser.add_argument("--temperature",
|
||||
type=float,
|
||||
default=0.8,
|
||||
help="Float that controls the randomness of the sampling.")
|
||||
parser.add_argument("--model-weight-gib",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Model weight memory usage in GiB (e.g., 1.0 for 0.5B model).")
|
||||
|
||||
args = parser.parse_args()
|
||||
if args.enable_sleep_mode:
|
||||
if args.model_weight_gib is None or args.temperature != 0:
|
||||
parser.error("model-weight-gib must be provided, and temperature must be zero when enable-sleep-mode is set.")
|
||||
if args.model_weight_gib <= 0:
|
||||
parser.error("model-weight-gib must be greater than 0 when enable-sleep-mode is set.")
|
||||
if args.model == parser.get_default("model") and args.model_weight_gib is None:
|
||||
parser.error("model-weight-gib must be provided for default model when enable-sleep-mode is set.")
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def main(
|
||||
local_rank: int,
|
||||
rank: int,
|
||||
master_addr: str,
|
||||
master_port: int,
|
||||
model_weight_gib: float,
|
||||
model: str = "Qwen/Qwen3-0.6B",
|
||||
world_size: int = 4,
|
||||
tensor_parallel_size: int = 2,
|
||||
enable_expert_parallel: bool = False,
|
||||
enforce_eager: bool = False,
|
||||
trust_remote_code: bool = True,
|
||||
enable_sleep_mode: bool = False,
|
||||
temperature: float = 0.8,
|
||||
):
|
||||
os.environ["MASTER_ADDR"] = master_addr
|
||||
os.environ["MASTER_PORT"] = str(master_port)
|
||||
os.environ["RANK"] = str(rank)
|
||||
os.environ["LOCAL_RANK"] = str(local_rank)
|
||||
os.environ["WORLD_SIZE"] = str(world_size)
|
||||
if not torch.distributed.is_initialized():
|
||||
torch.distributed.init_process_group(
|
||||
backend="cpu:gloo,npu:hccl",
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
)
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
] * 10
|
||||
sampling_params = SamplingParams(
|
||||
temperature=temperature,
|
||||
top_p=0.95,
|
||||
max_tokens=10,
|
||||
)
|
||||
llm = LLM(
|
||||
model=model,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
enable_expert_parallel=enable_expert_parallel,
|
||||
enforce_eager=enforce_eager,
|
||||
trust_remote_code=trust_remote_code,
|
||||
distributed_executor_backend="external_launcher",
|
||||
seed=0,
|
||||
enable_sleep_mode=enable_sleep_mode,
|
||||
)
|
||||
tp_ranks = get_tp_group().ranks
|
||||
print(f'TP RANKS: {tp_ranks}')
|
||||
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
if enable_sleep_mode:
|
||||
if rank == 0:
|
||||
free_bytes_before_sleep, total = torch.npu.mem_get_info()
|
||||
llm.sleep(level=1)
|
||||
if rank == 0:
|
||||
free_bytes_after_sleep, total = torch.npu.mem_get_info()
|
||||
freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
|
||||
print(f"Freed memory: {freed_bytes / 1024 ** 3:.2f} GiB")
|
||||
# now the freed memory should be larger than the model weights
|
||||
assert freed_bytes >= model_weight_gib / tensor_parallel_size * GiB_bytes
|
||||
|
||||
llm.wake_up()
|
||||
outputs_after_wakeup = llm.generate(prompts, sampling_params)
|
||||
if rank == 0:
|
||||
# cmp output
|
||||
assert outputs[0].outputs[0].text == outputs_after_wakeup[0].outputs[0].text
|
||||
print("Sleep and wake up successfully!!")
|
||||
|
||||
for i, output in enumerate(outputs):
|
||||
if i >= 5:
|
||||
# print only 5 outputs
|
||||
break
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Global rank: {rank}, Prompt: {prompt!r}, "
|
||||
f"Generated text: {generated_text!r}")
|
||||
|
||||
# Give engines time to pause their processing loops before exiting.
|
||||
sleep(5)
|
||||
del llm
|
||||
cleanup_env_and_memory()
|
||||
|
||||
|
||||
def cleanup_env_and_memory():
|
||||
destroy_model_parallel()
|
||||
destroy_distributed_environment()
|
||||
with contextlib.suppress(AssertionError):
|
||||
torch.distributed.destroy_process_group()
|
||||
gc.collect()
|
||||
torch.npu.empty_cache()
|
||||
torch.npu.reset_peak_memory_stats()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
|
||||
tp_size = args.tp_size
|
||||
node_size = args.node_size
|
||||
proc_per_node = args.proc_per_node
|
||||
node_rank = args.node_rank
|
||||
|
||||
if node_size == 1:
|
||||
master_addr = "127.0.0.1"
|
||||
master_port = get_open_port()
|
||||
else:
|
||||
master_addr = args.master_addr
|
||||
master_port = args.master_port
|
||||
|
||||
world_size = node_size * proc_per_node
|
||||
|
||||
procs = []
|
||||
for local_rank, rank in enumerate(
|
||||
range(proc_per_node * node_rank, proc_per_node * (node_rank + 1))):
|
||||
proc = Process(target=main,
|
||||
args=(
|
||||
local_rank,
|
||||
rank,
|
||||
master_addr,
|
||||
master_port,
|
||||
args.model_weight_gib,
|
||||
args.model,
|
||||
world_size,
|
||||
tp_size,
|
||||
args.enable_expert_parallel,
|
||||
args.enforce_eager,
|
||||
args.trust_remote_code,
|
||||
args.enable_sleep_mode,
|
||||
args.temperature,
|
||||
))
|
||||
|
||||
proc.start()
|
||||
procs.append(proc)
|
||||
exit_code = 0
|
||||
for proc in procs:
|
||||
proc.join(timeout=600)
|
||||
if proc.exitcode is None:
|
||||
print(
|
||||
f"Killing process {proc.pid} that didn't stop within 30 minutes."
|
||||
)
|
||||
proc.kill()
|
||||
exit_code = 1
|
||||
elif proc.exitcode:
|
||||
exit_code = proc.exitcode
|
||||
|
||||
exit(exit_code)
|
||||
105
examples/offline_inference_audio_language.py
Normal file
105
examples/offline_inference_audio_language.py
Normal file
@@ -0,0 +1,105 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
# Adapted from vllm-project/vllm/examples/offline_inference/audio_language.py
|
||||
#
|
||||
"""
|
||||
This example shows how to use vLLM for running offline inference
|
||||
with the correct prompt format on audio language models.
|
||||
|
||||
For most models, the prompt format should follow corresponding examples
|
||||
on HuggingFace model repository.
|
||||
"""
|
||||
|
||||
import os
|
||||
import argparse
|
||||
|
||||
from vllm.assets.audio import AudioAsset
|
||||
try:
|
||||
import librosa # type: ignore
|
||||
except ImportError:
|
||||
raise Exception("Can't import librosa, please ensure it's installed")
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
|
||||
def prepare_inputs(audio_count: int, audio_path1: str, audio_path2: str):
|
||||
use_vllm_audio_assert = True if audio_path1 == "mary_had_lamb" and audio_path2 == "winning_call" else False
|
||||
if use_vllm_audio_assert:
|
||||
audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
|
||||
else:
|
||||
audio_assets = [librosa.load(audio_path1, sr=None), librosa.load(audio_path2, sr=None)]
|
||||
|
||||
question_per_audio_count = {
|
||||
1: "What is recited in the audio?",
|
||||
2: "What sport and what nursery rhyme are referenced?"
|
||||
}
|
||||
|
||||
audio_in_prompt = "".join([
|
||||
f"Audio {idx+1}: <|audio_bos|><|AUDIO|><|audio_eos|>\n"
|
||||
for idx in range(audio_count)
|
||||
])
|
||||
question = question_per_audio_count[audio_count]
|
||||
prompt = ("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
f"{audio_in_prompt}{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n")
|
||||
|
||||
mm_data = {
|
||||
"audio":
|
||||
audio_assets if not use_vllm_audio_assert else [asset.audio_and_sample_rate for asset in audio_assets[:audio_count]]
|
||||
}
|
||||
|
||||
# Merge text prompt and audio data into inputs
|
||||
inputs = {"prompt": prompt, "multi_modal_data": mm_data}
|
||||
return inputs
|
||||
|
||||
|
||||
def main(audio_count: int, audio_path1: str, audio_path2: str):
|
||||
# NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on
|
||||
# lower-end GPUs.
|
||||
# Unless specified, these settings have been tested to work on a single L4.
|
||||
# `limit_mm_per_prompt`: the max num items for each modality per prompt.
|
||||
llm = LLM(model="Qwen/Qwen2-Audio-7B-Instruct",
|
||||
max_model_len=4096,
|
||||
max_num_seqs=5,
|
||||
limit_mm_per_prompt={"audio": audio_count},
|
||||
enforce_eager=True)
|
||||
|
||||
inputs = prepare_inputs(audio_count, audio_path1, audio_path2)
|
||||
|
||||
sampling_params = SamplingParams(temperature=0.2,
|
||||
max_tokens=64,
|
||||
stop_token_ids=None)
|
||||
|
||||
outputs = llm.generate(inputs, sampling_params=sampling_params)
|
||||
|
||||
for o in outputs:
|
||||
generated_text = o.outputs[0].text
|
||||
print("generated_text:", generated_text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Arguments of rank table generator", )
|
||||
parser.add_argument("--audio-path1", type=str, default="mary_had_lamb")
|
||||
parser.add_argument("--audio-path2", type=str, default="winning_call")
|
||||
args = parser.parse_args()
|
||||
|
||||
audio_count = 2
|
||||
main(audio_count, args.audio_path1, args.audio_path2)
|
||||
51
examples/offline_inference_npu.py
Normal file
51
examples/offline_inference_npu.py
Normal file
@@ -0,0 +1,51 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
# Adapted from vllm-project/vllm/examples/offline_inference/basic.py
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
# isort: skip_file
|
||||
import os
|
||||
|
||||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
|
||||
def main():
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
# Create a sampling params object.
|
||||
sampling_params = SamplingParams(max_tokens=100, temperature=0.0)
|
||||
# Create an LLM.
|
||||
llm = LLM(model="Qwen/Qwen2.5-0.5B-Instruct")
|
||||
|
||||
# Generate texts from the prompts.
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
55
examples/offline_inference_npu_tp2.py
Normal file
55
examples/offline_inference_npu_tp2.py
Normal file
@@ -0,0 +1,55 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
# Adapted from vllm-project/vllm/examples/offline_inference/basic.py
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
# isort: skip_file
|
||||
import os
|
||||
|
||||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
|
||||
def main():
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
# Create a sampling params object.
|
||||
sampling_params = SamplingParams(max_tokens=100, temperature=0.0)
|
||||
# Create an LLM.
|
||||
llm = LLM(model="deepseek-ai/DeepSeek-V2-Lite",
|
||||
tensor_parallel_size=2,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
max_model_len=1024)
|
||||
|
||||
# Generate texts from the prompts.
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
57
examples/offline_inference_sleep_mode_npu.py
Normal file
57
examples/offline_inference_sleep_mode_npu.py
Normal file
@@ -0,0 +1,57 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.utils import GiB_bytes
|
||||
|
||||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
def main():
|
||||
prompt = "How are you?"
|
||||
|
||||
free, total = torch.npu.mem_get_info()
|
||||
print(f"Free memory before sleep: {free / 1024 ** 3:.2f} GiB")
|
||||
# record npu memory use baseline in case other process is running
|
||||
used_bytes_baseline = total - free
|
||||
llm = LLM("Qwen/Qwen2.5-0.5B-Instruct", enable_sleep_mode=True)
|
||||
sampling_params = SamplingParams(temperature=0, max_tokens=10)
|
||||
output = llm.generate(prompt, sampling_params)
|
||||
|
||||
llm.sleep(level=1)
|
||||
|
||||
free_npu_bytes_after_sleep, total = torch.npu.mem_get_info()
|
||||
print(
|
||||
f"Free memory after sleep: {free_npu_bytes_after_sleep / 1024 ** 3:.2f} GiB"
|
||||
)
|
||||
used_bytes = total - free_npu_bytes_after_sleep - used_bytes_baseline
|
||||
# now the memory usage should be less than the model weights
|
||||
# (0.5B model, 1GiB weights)
|
||||
assert used_bytes < 1 * GiB_bytes
|
||||
|
||||
llm.wake_up()
|
||||
output2 = llm.generate(prompt, sampling_params)
|
||||
# cmp output
|
||||
assert output[0].outputs[0].text == output2[0].outputs[0].text
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
88
examples/prompt_embedding_inference.py
Normal file
88
examples/prompt_embedding_inference.py
Normal file
@@ -0,0 +1,88 @@
|
||||
import os
|
||||
|
||||
import torch
|
||||
from transformers import (AutoModelForCausalLM, AutoTokenizer,
|
||||
PreTrainedTokenizer)
|
||||
from vllm import LLM
|
||||
|
||||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
|
||||
def init_tokenizer_and_llm(model_name: str):
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
transformers_model = AutoModelForCausalLM.from_pretrained(model_name)
|
||||
embedding_layer = transformers_model.get_input_embeddings()
|
||||
llm = LLM(model=model_name, enable_prompt_embeds=True)
|
||||
return tokenizer, embedding_layer, llm
|
||||
|
||||
|
||||
def get_prompt_embeds(chat: list[dict[str,
|
||||
str]], tokenizer: PreTrainedTokenizer,
|
||||
embedding_layer: torch.nn.Module):
|
||||
token_ids = tokenizer.apply_chat_template(chat,
|
||||
add_generation_prompt=True,
|
||||
return_tensors='pt')
|
||||
prompt_embeds = embedding_layer(token_ids).squeeze(0)
|
||||
return prompt_embeds
|
||||
|
||||
|
||||
def single_prompt_inference(llm: LLM, tokenizer: PreTrainedTokenizer,
|
||||
embedding_layer: torch.nn.Module):
|
||||
chat = [{
|
||||
"role": "user",
|
||||
"content": "Please tell me about the capital of France."
|
||||
}]
|
||||
prompt_embeds = get_prompt_embeds(chat, tokenizer, embedding_layer)
|
||||
|
||||
outputs = llm.generate({
|
||||
"prompt_embeds": prompt_embeds,
|
||||
})
|
||||
|
||||
print("\n[Single Inference Output]")
|
||||
print("-" * 30)
|
||||
for o in outputs:
|
||||
print(o.outputs[0].text)
|
||||
print("-" * 30)
|
||||
|
||||
|
||||
def batch_prompt_inference(llm: LLM, tokenizer: PreTrainedTokenizer,
|
||||
embedding_layer: torch.nn.Module):
|
||||
chats = [[{
|
||||
"role": "user",
|
||||
"content": "Please tell me about the capital of France."
|
||||
}],
|
||||
[{
|
||||
"role": "user",
|
||||
"content": "When is the day longest during the year?"
|
||||
}],
|
||||
[{
|
||||
"role": "user",
|
||||
"content": "Where is bigger, the moon or the sun?"
|
||||
}]]
|
||||
|
||||
prompt_embeds_list = [
|
||||
get_prompt_embeds(chat, tokenizer, embedding_layer) for chat in chats
|
||||
]
|
||||
|
||||
outputs = llm.generate([{
|
||||
"prompt_embeds": embeds
|
||||
} for embeds in prompt_embeds_list])
|
||||
|
||||
print("\n[Batch Inference Outputs]")
|
||||
print("-" * 30)
|
||||
for i, o in enumerate(outputs):
|
||||
print(f"Q{i+1}: {chats[i][0]['content']}")
|
||||
print(f"A{i+1}: {o.outputs[0].text}\n")
|
||||
print("-" * 30)
|
||||
|
||||
|
||||
def main():
|
||||
model_name = "meta-llama/Llama-3.2-1B-Instruct"
|
||||
tokenizer, embedding_layer, llm = init_tokenizer_and_llm(model_name)
|
||||
single_prompt_inference(llm, tokenizer, embedding_layer)
|
||||
batch_prompt_inference(llm, tokenizer, embedding_layer)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
32
examples/run_dp_server.sh
Normal file
32
examples/run_dp_server.sh
Normal file
@@ -0,0 +1,32 @@
|
||||
|
||||
export HCCL_IF_IP=2.0.0.0
|
||||
export GLOO_SOCKET_IFNAME="eth0"
|
||||
export TP_SOCKET_IFNAME="eth0"
|
||||
export HCCL_SOCKET_IFNAME="eth0"
|
||||
|
||||
export OMP_PROC_BIND=false
|
||||
export OMP_NUM_THREADS=100
|
||||
|
||||
export VLLM_USE_V1=1
|
||||
export VLLM_USE_MODELSCOPE=true
|
||||
|
||||
export ASCEND_LAUNCH_BLOCKING=0
|
||||
|
||||
vllm serve Qwen/Qwen1.5-MoE-A2.7B \
|
||||
--host 0.0.0.0 \
|
||||
--port 20002 \
|
||||
--served-model-name Qwen \
|
||||
--data-parallel-size 2 \
|
||||
--data-parallel-size-local 2 \
|
||||
--data-parallel-address 2.0.0.0 \
|
||||
--data-parallel-rpc-port 13389 \
|
||||
--tensor-parallel-size 4 \
|
||||
--enable-expert-parallel \
|
||||
--no-enable-prefix-caching \
|
||||
--max-num-seqs 16 \
|
||||
--max-model-len 4096 \
|
||||
--max-num-batched-tokens 4096 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--additional-config '{"ascend_scheduler_config":{"enabled":true},"torchair_graph_config":{"enabled":false, "enable_multistream_moe":false, "use_cached_graph":false}}'
|
||||
Reference in New Issue
Block a user