193 lines
6.9 KiB
Python
193 lines
6.9 KiB
Python
#!/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("--model", 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
|
|
from vllm.config import (
|
|
ModelConfig,
|
|
ParallelConfig,
|
|
VllmConfig,
|
|
set_current_vllm_config,
|
|
)
|
|
|
|
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()
|
|
model_config = ModelConfig(
|
|
model=str(args.model),
|
|
dtype="bfloat16",
|
|
max_model_len=8192,
|
|
skip_tokenizer_init=True,
|
|
generation_config="vllm",
|
|
)
|
|
vllm_config = VllmConfig(
|
|
model_config=model_config,
|
|
parallel_config=ParallelConfig(tensor_parallel_size=world_size)
|
|
)
|
|
with set_current_vllm_config(vllm_config):
|
|
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}"
|
|
)
|
|
uses_flashinfer = fi_comm.should_use_fi_ar(input_tensor)
|
|
|
|
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" if uses_flashinfer else "nccl_fallback"
|
|
),
|
|
"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),
|
|
"model": str(args.model),
|
|
"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()
|