Add FlashInfer TRTLLM all-reduce profile smoke
This commit is contained in:
@@ -0,0 +1,19 @@
|
||||
version = 1
|
||||
|
||||
[[jobs]]
|
||||
name = "qwen30-vllm020-flashinfer-allreduce-smoke-tp2-20260716-v1"
|
||||
gpus = 2
|
||||
gpu_model = "H20"
|
||||
hosts = ["dash0"]
|
||||
command = "cd /home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-v1/runs/frontier-qwen30-vllm020-profile-v1 && timeout --signal=TERM --kill-after=30s 840 bash run_allreduce_profile.sh"
|
||||
artifacts = ["artifacts/allreduce-smoke-tp2"]
|
||||
|
||||
[jobs.env]
|
||||
HOME = "/tmp/wjh"
|
||||
XDG_CACHE_HOME = "/tmp/wjh/.cache"
|
||||
VLLM_CACHE_ROOT = "/tmp/wjh/.cache/vllm"
|
||||
TP = "2"
|
||||
NUM_TOKENS = "8"
|
||||
OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/allreduce-smoke-tp2"
|
||||
VENV_ROOT = "/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1"
|
||||
VLLM_SOURCE = "/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build"
|
||||
@@ -0,0 +1,173 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Profile vLLM 0.20 TP all-reduce and assert FlashInfer TRTLLM dispatch."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import statistics
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import vllm
|
||||
|
||||
|
||||
VLLM_VERSION = "0.20.0"
|
||||
VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1"
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--vllm-source", type=Path, required=True)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--num-tokens", type=int, nargs="+", default=[8])
|
||||
parser.add_argument("--hidden-dim", type=int, default=2048)
|
||||
parser.add_argument("--warmup-iters", type=int, default=3)
|
||||
parser.add_argument("--repeats", type=int, default=10)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def git_head(repo: Path) -> str:
|
||||
return subprocess.check_output(
|
||||
["git", "-C", str(repo), "rev-parse", "HEAD"], text=True
|
||||
).strip()
|
||||
|
||||
|
||||
def stats_ms(samples: list[float]) -> dict[str, float]:
|
||||
return {
|
||||
"min": min(samples),
|
||||
"max": max(samples),
|
||||
"mean": statistics.fmean(samples),
|
||||
"median": statistics.median(samples),
|
||||
"std": statistics.pstdev(samples),
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if vllm.__version__ != VLLM_VERSION:
|
||||
raise SystemExit(f"expected vLLM {VLLM_VERSION}, got {vllm.__version__}")
|
||||
source_head = git_head(args.vllm_source)
|
||||
if source_head != VLLM_COMMIT:
|
||||
raise SystemExit(f"expected vLLM source {VLLM_COMMIT}, got {source_head}")
|
||||
if os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER") != "1":
|
||||
raise SystemExit("VLLM_ALLREDUCE_USE_FLASHINFER must equal 1")
|
||||
if os.getenv("VLLM_FLASHINFER_ALLREDUCE_BACKEND") != "trtllm":
|
||||
raise SystemExit("VLLM_FLASHINFER_ALLREDUCE_BACKEND must equal trtllm")
|
||||
if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ:
|
||||
raise SystemExit("launch with torchrun")
|
||||
|
||||
from vllm.distributed import (
|
||||
destroy_distributed_environment,
|
||||
destroy_model_parallel,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from vllm.distributed.parallel_state import get_tp_group
|
||||
|
||||
rank = int(os.environ["RANK"])
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
if world_size not in (2, 4):
|
||||
raise SystemExit(f"expected TP world size 2 or 4, got {world_size}")
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
torch.set_default_device(device)
|
||||
init_distributed_environment()
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
rows: list[dict[str, Any]] = []
|
||||
expected_sum = world_size * (world_size + 1) / 2
|
||||
try:
|
||||
for num_tokens in args.num_tokens:
|
||||
input_tensor = torch.full(
|
||||
(num_tokens, args.hidden_dim),
|
||||
float(rank + 1),
|
||||
dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
for _ in range(args.warmup_iters):
|
||||
output = tensor_model_parallel_all_reduce(input_tensor)
|
||||
torch.accelerator.synchronize()
|
||||
torch.testing.assert_close(
|
||||
output,
|
||||
torch.full_like(output, expected_sum),
|
||||
atol=0.0,
|
||||
rtol=0.0,
|
||||
)
|
||||
|
||||
communicator = get_tp_group().device_communicator
|
||||
fi_comm = communicator.fi_ar_comm
|
||||
if fi_comm is None or fi_comm.disabled:
|
||||
raise SystemExit(
|
||||
f"FlashInfer all-reduce was not selected at TP={world_size}, "
|
||||
f"tokens={num_tokens}"
|
||||
)
|
||||
if not fi_comm.should_use_fi_ar(input_tensor):
|
||||
raise SystemExit(
|
||||
f"FlashInfer rejected serving tensor shape {tuple(input_tensor.shape)}"
|
||||
)
|
||||
|
||||
samples: list[float] = []
|
||||
for _ in range(args.repeats):
|
||||
dist.barrier()
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
output = tensor_model_parallel_all_reduce(input_tensor)
|
||||
end.record()
|
||||
torch.accelerator.synchronize()
|
||||
samples.append(float(start.elapsed_time(end)))
|
||||
|
||||
gathered: list[list[float] | None] = [None] * world_size
|
||||
dist.all_gather_object(gathered, samples)
|
||||
if rank == 0:
|
||||
per_rank = [stats_ms(item) for item in gathered if item is not None]
|
||||
row = {
|
||||
"tensor_parallel_size": world_size,
|
||||
"num_tokens": num_tokens,
|
||||
"hidden_dim": args.hidden_dim,
|
||||
"payload_bytes": num_tokens
|
||||
* args.hidden_dim
|
||||
* torch.tensor([], dtype=torch.bfloat16).element_size(),
|
||||
"dtype": "bfloat16",
|
||||
"communicator": "vllm.tensor_model_parallel_all_reduce",
|
||||
"selected_backend": "flashinfer_trtllm",
|
||||
"per_rank_time_ms": per_rank,
|
||||
"critical_path_median_ms": max(
|
||||
rank_stats["median"] for rank_stats in per_rank
|
||||
),
|
||||
}
|
||||
rows.append(row)
|
||||
print(json.dumps(row, sort_keys=True), flush=True)
|
||||
finally:
|
||||
destroy_model_parallel()
|
||||
destroy_distributed_environment()
|
||||
|
||||
if rank == 0:
|
||||
payload = {
|
||||
"schema_version": "qwen30_vllm020_allreduce_raw.v1",
|
||||
"environment": {
|
||||
"vllm_version": vllm.__version__,
|
||||
"vllm_source_commit": source_head,
|
||||
"torch_version": torch.__version__,
|
||||
"torch_cuda": torch.version.cuda,
|
||||
"gpu": torch.cuda.get_device_name(device),
|
||||
"backend_env": {
|
||||
"VLLM_ALLREDUCE_USE_FLASHINFER": "1",
|
||||
"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "trtllm",
|
||||
},
|
||||
},
|
||||
"rows": rows,
|
||||
}
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,53 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
TP="${TP:?TP must be set to 2 or 4}"
|
||||
case "${TP}" in
|
||||
2) HARD_GPU_CAP="0.40_H20h" ;;
|
||||
4) HARD_GPU_CAP="0.80_H20h" ;;
|
||||
*) echo "ERROR: invalid TP=${TP}" >&2; exit 1 ;;
|
||||
esac
|
||||
OUTPUT_ROOT="${OUTPUT_ROOT:?OUTPUT_ROOT must be set}"
|
||||
VENV_ROOT="${VENV_ROOT:-/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1}"
|
||||
VLLM_SOURCE="${VLLM_SOURCE:-/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build}"
|
||||
NUM_TOKENS="${NUM_TOKENS:-8}"
|
||||
mkdir -p "${OUTPUT_ROOT}/logs" "${OUTPUT_ROOT}/provenance" "${OUTPUT_ROOT}/raw"
|
||||
exec > >(tee -a "${OUTPUT_ROOT}/logs/profile.log") 2>&1
|
||||
|
||||
IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES:?fleet GPUs are required}"
|
||||
if [[ "${#GPU_IDS[@]}" -ne "${TP}" ]]; then
|
||||
echo "ERROR: TP=${TP} requires ${TP} GPUs, got ${CUDA_VISIBLE_DEVICES}" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
export VLLM_ALLREDUCE_USE_FLASHINFER=1
|
||||
export VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm
|
||||
echo "PROFILE_LAUNCH_ECHO host=$(hostname) gpus=${CUDA_VISIBLE_DEVICES} runtime=vLLM-0.20.0+cu129 operator=tensor_model_parallel_all_reduce backend=FlashInfer-TRTLLM tp=${TP} tokens=${NUM_TOKENS} hidden=2048 dtype=BF16 output=${OUTPUT_ROOT} expected_wall=2-6m hard_wall=720s hard_gpu_cap=${HARD_GPU_CAP}"
|
||||
date -u +"START_UTC=%Y-%m-%dT%H:%M:%SZ"
|
||||
nvidia-smi --query-gpu=index,name,driver_version,memory.used,utilization.gpu --format=csv,noheader
|
||||
|
||||
git rev-parse HEAD > "${OUTPUT_ROOT}/provenance/aituner.commit"
|
||||
git -C "${VLLM_SOURCE}" rev-parse HEAD > "${OUTPUT_ROOT}/provenance/vllm-source.commit"
|
||||
sha256sum profile_vllm020_allreduce.py run_allreduce_profile.sh \
|
||||
> "${OUTPUT_ROOT}/provenance/source.sha256"
|
||||
uv pip freeze --python "${VENV_ROOT}/bin/python" \
|
||||
> "${OUTPUT_ROOT}/provenance/pip-freeze.txt"
|
||||
nvidia-smi --query-gpu=index,uuid,name,driver_version,memory.total \
|
||||
--format=csv,noheader > "${OUTPUT_ROOT}/provenance/gpus.csv"
|
||||
|
||||
read -r -a TOKEN_ARGS <<< "${NUM_TOKENS}"
|
||||
timeout --signal=TERM --kill-after=30s 600 \
|
||||
"${VENV_ROOT}/bin/torchrun" --standalone --nproc_per_node="${TP}" \
|
||||
profile_vllm020_allreduce.py \
|
||||
--vllm-source "${VLLM_SOURCE}" \
|
||||
--output "${OUTPUT_ROOT}/raw/allreduce-tp${TP}.json" \
|
||||
--num-tokens "${TOKEN_ARGS[@]}" \
|
||||
--warmup-iters 3 \
|
||||
--repeats 10
|
||||
|
||||
test -s "${OUTPUT_ROOT}/raw/allreduce-tp${TP}.json"
|
||||
sha256sum "${OUTPUT_ROOT}/raw/allreduce-tp${TP}.json" \
|
||||
"${OUTPUT_ROOT}/provenance"/* > "${OUTPUT_ROOT}/artifacts.sha256"
|
||||
date -u +"END_UTC=%Y-%m-%dT%H:%M:%SZ"
|
||||
echo "ALLREDUCE_PROFILE_COMPLETE tp=${TP} tokens=${NUM_TOKENS}"
|
||||
Reference in New Issue
Block a user