diff --git a/runs/frontier-qwen30-vllm020-profile-v1/jobs_allreduce_smoke_tp2.toml b/runs/frontier-qwen30-vllm020-profile-v1/jobs_allreduce_smoke_tp2.toml new file mode 100644 index 0000000..802066e --- /dev/null +++ b/runs/frontier-qwen30-vllm020-profile-v1/jobs_allreduce_smoke_tp2.toml @@ -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" diff --git a/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_allreduce.py b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_allreduce.py new file mode 100644 index 0000000..8665b9c --- /dev/null +++ b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_allreduce.py @@ -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() diff --git a/runs/frontier-qwen30-vllm020-profile-v1/run_allreduce_profile.sh b/runs/frontier-qwen30-vllm020-profile-v1/run_allreduce_profile.sh new file mode 100644 index 0000000..6321a8d --- /dev/null +++ b/runs/frontier-qwen30-vllm020-profile-v1/run_allreduce_profile.sh @@ -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}"