346 lines
14 KiB
Python
346 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Profile the vLLM 0.20 TP all-reduce path under a frozen runtime contract."""
|
|
|
|
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(
|
|
"--collective-contract",
|
|
choices=("flashinfer-trtllm", "qwen235-serving-projected"),
|
|
default="flashinfer-trtllm",
|
|
)
|
|
parser.add_argument("--warmup-iters", type=int, default=3)
|
|
parser.add_argument("--repeats", type=int, default=10)
|
|
parser.add_argument("--trials", type=int, default=1)
|
|
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 args.repeats <= 0 or args.trials <= 0 or args.warmup_iters < 0:
|
|
raise SystemExit("warmup/repeats/trials must be non-negative/positive")
|
|
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")
|
|
expected_symm_mem = (
|
|
"1" if args.collective_contract == "qwen235-serving-projected" else "0"
|
|
)
|
|
if os.getenv("VLLM_ALLREDUCE_USE_SYMM_MEM") != expected_symm_mem:
|
|
raise SystemExit(
|
|
f"VLLM_ALLREDUCE_USE_SYMM_MEM must equal {expected_symm_mem}"
|
|
)
|
|
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,
|
|
set_custom_all_reduce,
|
|
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, 8):
|
|
raise SystemExit(f"expected TP world size 2, 4, or 8, got {world_size}")
|
|
device = torch.device(f"cuda:{local_rank}")
|
|
torch.accelerator.set_device_index(device)
|
|
torch.set_default_device(device)
|
|
disable_custom_all_reduce = (
|
|
args.collective_contract == "qwen235-serving-projected"
|
|
)
|
|
set_custom_all_reduce(not disable_custom_all_reduce)
|
|
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,
|
|
disable_custom_all_reduce=disable_custom_all_reduce,
|
|
)
|
|
)
|
|
with set_current_vllm_config(vllm_config):
|
|
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
|
|
|
from vllm.distributed.device_communicators.all_reduce_utils import (
|
|
should_nccl_symm_mem_allreduce,
|
|
)
|
|
|
|
def resolve_profile_operation(input_tensor: torch.Tensor):
|
|
communicator = get_tp_group().device_communicator
|
|
if args.collective_contract == "qwen235-serving-projected":
|
|
from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
|
|
FI_ALLREDUCE_FUSION_MAX_SIZE_MB,
|
|
)
|
|
|
|
major, minor = torch.cuda.get_device_capability(device)
|
|
capability = major * 10 + minor
|
|
try:
|
|
fusion_limit_bytes = int(
|
|
FI_ALLREDUCE_FUSION_MAX_SIZE_MB[capability][world_size]
|
|
* 1024
|
|
* 1024
|
|
)
|
|
except KeyError as error:
|
|
raise SystemExit(
|
|
f"no FlashInfer fusion limit for SM{capability}, TP{world_size}"
|
|
) from error
|
|
if input_tensor.nbytes <= fusion_limit_bytes:
|
|
from vllm.distributed.device_communicators.flashinfer_all_reduce import (
|
|
get_fi_ar_workspace,
|
|
)
|
|
|
|
fi_comm = communicator.fi_ar_comm
|
|
if fi_comm is None or fi_comm.disabled:
|
|
raise SystemExit("FlashInfer all-reduce communicator is unavailable")
|
|
max_token_num = fusion_limit_bytes // (
|
|
args.hidden_dim * input_tensor.element_size()
|
|
)
|
|
workspace = get_fi_ar_workspace(
|
|
world_size=world_size,
|
|
rank=rank,
|
|
max_token_num=max_token_num,
|
|
hidden_dim=args.hidden_dim,
|
|
dtype=input_tensor.dtype,
|
|
group=get_tp_group().device_group,
|
|
)
|
|
if workspace is None:
|
|
raise SystemExit(
|
|
"FlashInfer rejected the serving-matched integer workspace"
|
|
)
|
|
if not fi_comm.should_use_fi_ar(input_tensor):
|
|
raise SystemExit("FlashInfer rejected a fusion-eligible payload")
|
|
return (
|
|
"flashinfer_trtllm_fused_projection",
|
|
fi_comm.all_reduce,
|
|
fusion_limit_bytes,
|
|
)
|
|
if (
|
|
communicator.pynccl_comm is not None
|
|
and should_nccl_symm_mem_allreduce(world_size, input_tensor)
|
|
):
|
|
return (
|
|
"pynccl_symmetric_with_copy",
|
|
torch.ops.vllm.all_reduce_symmetric_with_copy,
|
|
fusion_limit_bytes,
|
|
)
|
|
if (
|
|
communicator.symm_mem_comm is not None
|
|
and communicator.symm_mem_comm.should_use_symm_mem(input_tensor)
|
|
):
|
|
return (
|
|
"torch_symmetric_memory",
|
|
communicator.symm_mem_comm.all_reduce,
|
|
fusion_limit_bytes,
|
|
)
|
|
if communicator.pynccl_comm is not None:
|
|
return "pynccl", communicator.pynccl_comm.all_reduce, fusion_limit_bytes
|
|
return "torch_distributed", tensor_model_parallel_all_reduce, fusion_limit_bytes
|
|
|
|
fi_comm = communicator.fi_ar_comm
|
|
selected_backend = (
|
|
"flashinfer_trtllm"
|
|
if fi_comm is not None
|
|
and not fi_comm.disabled
|
|
and fi_comm.should_use_fi_ar(input_tensor)
|
|
else "non_flashinfer_fallback"
|
|
)
|
|
return selected_backend, tensor_model_parallel_all_reduce, None
|
|
|
|
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,
|
|
)
|
|
selected_backend, profile_operation, fusion_limit_bytes = (
|
|
resolve_profile_operation(input_tensor)
|
|
)
|
|
for _ in range(args.warmup_iters):
|
|
output = profile_operation(input_tensor)
|
|
torch.accelerator.synchronize()
|
|
torch.testing.assert_close(
|
|
output,
|
|
torch.full_like(output, expected_sum),
|
|
atol=0.0,
|
|
rtol=0.0,
|
|
)
|
|
|
|
if args.collective_contract == "flashinfer-trtllm":
|
|
if selected_backend != "flashinfer_trtllm":
|
|
raise SystemExit(
|
|
f"expected FlashInfer TRTLLM, got {selected_backend} at "
|
|
f"TP={world_size}, tokens={num_tokens}"
|
|
)
|
|
elif selected_backend not in {
|
|
"flashinfer_trtllm_fused_projection",
|
|
"pynccl_symmetric_with_copy",
|
|
"torch_symmetric_memory",
|
|
"pynccl",
|
|
}:
|
|
raise SystemExit(
|
|
f"unexpected Qwen235 serving backend: {selected_backend}"
|
|
)
|
|
|
|
all_trial_rank_samples: list[list[list[float]]] = []
|
|
critical_path_samples: list[float] = []
|
|
for _ in range(args.trials):
|
|
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 = profile_operation(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:
|
|
rank_samples = [item for item in gathered if item is not None]
|
|
all_trial_rank_samples.append(rank_samples)
|
|
critical_path_samples.extend(
|
|
max(per_rank[index] for per_rank in rank_samples)
|
|
for index in range(args.repeats)
|
|
)
|
|
if rank == 0:
|
|
flattened_by_rank = [
|
|
[
|
|
sample
|
|
for trial in all_trial_rank_samples
|
|
for sample in trial[rank_index]
|
|
]
|
|
for rank_index in range(world_size)
|
|
]
|
|
per_rank = [stats_ms(samples) for samples in flattened_by_rank]
|
|
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",
|
|
"collective_contract": args.collective_contract,
|
|
"selected_backend": selected_backend,
|
|
"real_fusion_limit_bytes": fusion_limit_bytes,
|
|
"trials": args.trials,
|
|
"repeats_per_trial": args.repeats,
|
|
"per_trial_rank_samples_ms": all_trial_rank_samples,
|
|
"per_rank_time_ms": per_rank,
|
|
"critical_path_time_ms": stats_ms(critical_path_samples),
|
|
"critical_path_median_ms": statistics.median(
|
|
critical_path_samples
|
|
),
|
|
}
|
|
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"
|
|
if args.collective_contract == "flashinfer-trtllm"
|
|
else "vllm020_allreduce_raw.v2"
|
|
),
|
|
"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": os.getenv(
|
|
"VLLM_ALLREDUCE_USE_FLASHINFER", "0"
|
|
),
|
|
"VLLM_FLASHINFER_ALLREDUCE_BACKEND": os.getenv(
|
|
"VLLM_FLASHINFER_ALLREDUCE_BACKEND"
|
|
),
|
|
"VLLM_ALLREDUCE_USE_SYMM_MEM": os.getenv(
|
|
"VLLM_ALLREDUCE_USE_SYMM_MEM", "1"
|
|
),
|
|
},
|
|
"collective_contract": args.collective_contract,
|
|
"disable_custom_all_reduce": disable_custom_all_reduce,
|
|
},
|
|
"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()
|