From 07b1eb4b75ae73f6156ba9ee696c1b0083098317 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Thu, 16 Jul 2026 22:54:17 +0800 Subject: [PATCH] Initialize TP context for router profiling --- .../jobs_router_full.toml | 6 +- .../profile_vllm020_router.py | 138 +++++++++++------- 2 files changed, 85 insertions(+), 59 deletions(-) diff --git a/runs/frontier-qwen30-vllm020-profile-v1/jobs_router_full.toml b/runs/frontier-qwen30-vllm020-profile-v1/jobs_router_full.toml index 15719fc..4dd0908 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/jobs_router_full.toml +++ b/runs/frontier-qwen30-vllm020-profile-v1/jobs_router_full.toml @@ -1,18 +1,18 @@ version = 1 [[jobs]] -name = "qwen30-vllm020-router-full-20260716-v2-disable-tp-init" +name = "qwen30-vllm020-router-full-20260716-v3-tp-context" gpus = 1 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 1020 bash run_router_full.sh" -artifacts = ["artifacts/router-full-v2"] +artifacts = ["artifacts/router-full-v3"] [jobs.env] HOME = "/tmp/wjh" XDG_CACHE_HOME = "/tmp/wjh/.cache" VLLM_CACHE_ROOT = "/tmp/wjh/.cache/vllm" -OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/router-full-v2" +OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/router-full-v3" 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" MODEL = "/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B" diff --git a/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py index 7e23862..8da2070 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py +++ b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py @@ -5,6 +5,7 @@ from __future__ import annotations import argparse import json +import socket import statistics import subprocess from pathlib import Path @@ -93,6 +94,12 @@ def main() -> None: raise SystemExit(f"model contract mismatch: expected {expected}, got {observed}") from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config + from vllm.distributed import ( + destroy_distributed_environment, + destroy_model_parallel, + init_distributed_environment, + initialize_model_parallel, + ) from vllm.model_executor.layers.fused_moe import fused_topk from vllm.model_executor.layers.linear import ReplicatedLinear @@ -108,65 +115,84 @@ def main() -> None: ) rows: list[dict[str, Any]] = [] - with set_current_vllm_config(VllmConfig(model_config=model_config)): - gate = ReplicatedLinear( - HIDDEN_DIM, - NUM_EXPERTS, - bias=False, - quant_config=None, - prefix="model.layers.0.mlp.gate", - disable_tp=True, - ).to(device=device, dtype=torch.bfloat16) - gate.weight.data.uniform_(-0.01, 0.01) + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener: + listener.bind(("127.0.0.1", 0)) + distributed_init_method = f"tcp://127.0.0.1:{listener.getsockname()[1]}" + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=distributed_init_method, + ) + try: + with set_current_vllm_config(VllmConfig(model_config=model_config)): + initialize_model_parallel(tensor_model_parallel_size=1) + gate = ReplicatedLinear( + HIDDEN_DIM, + NUM_EXPERTS, + bias=False, + quant_config=None, + prefix="model.layers.0.mlp.gate", + ).to(device=device, dtype=torch.bfloat16) + gate.weight.data.uniform_(-0.01, 0.01) - for num_tokens in args.num_tokens: - hidden = torch.empty( - (num_tokens, HIDDEN_DIM), device=device, dtype=torch.bfloat16 - ).uniform_(-0.1, 0.1) - logits, gate_time = measure_ms( - lambda: gate(hidden)[0], args.warmup_iters, args.repeats - ) - topk_result, topk_time = measure_ms( - lambda: fused_topk(hidden, logits, TOP_K, renormalize=True), - args.warmup_iters, - args.repeats, - ) + for num_tokens in args.num_tokens: + hidden = torch.empty( + (num_tokens, HIDDEN_DIM), device=device, dtype=torch.bfloat16 + ).uniform_(-0.1, 0.1) + logits, gate_time = measure_ms( + lambda: gate(hidden)[0], args.warmup_iters, args.repeats + ) + topk_result, topk_time = measure_ms( + lambda: fused_topk(hidden, logits, TOP_K, renormalize=True), + args.warmup_iters, + args.repeats, + ) - def gate_and_topk() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - current_logits, _ = gate(hidden) - return fused_topk(hidden, current_logits, TOP_K, renormalize=True) + def gate_and_topk() -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor + ]: + current_logits, _ = gate(hidden) + return fused_topk( + hidden, current_logits, TOP_K, renormalize=True + ) - combined_result, combined_time = measure_ms( - gate_and_topk, args.warmup_iters, args.repeats - ) - topk_weights, topk_ids, _ = topk_result - combined_weights, combined_ids, _ = combined_result - if logits.shape != (num_tokens, NUM_EXPERTS): - raise SystemExit(f"invalid gate output shape: {tuple(logits.shape)}") - if topk_ids.shape != (num_tokens, TOP_K): - raise SystemExit(f"invalid top-k shape: {tuple(topk_ids.shape)}") - torch.testing.assert_close( - topk_weights.sum(dim=-1), - torch.ones(num_tokens, device=device), - atol=1e-5, - rtol=1e-5, - ) - torch.testing.assert_close(combined_weights, topk_weights) - torch.testing.assert_close(combined_ids, topk_ids) - additive_median = gate_time["median"] + topk_time["median"] - row = { - "num_tokens": num_tokens, - "gate_linear_time_ms": gate_time, - "routing_topk_time_ms": topk_time, - "gate_plus_topk_time_ms": combined_time, - "median_nonadditivity_ratio": ( - combined_time["median"] / additive_median - if additive_median > 0 - else 1.0 - ), - } - rows.append(row) - print(json.dumps(row, sort_keys=True), flush=True) + combined_result, combined_time = measure_ms( + gate_and_topk, args.warmup_iters, args.repeats + ) + topk_weights, topk_ids, _ = topk_result + combined_weights, combined_ids, _ = combined_result + if logits.shape != (num_tokens, NUM_EXPERTS): + raise SystemExit( + f"invalid gate output shape: {tuple(logits.shape)}" + ) + if topk_ids.shape != (num_tokens, TOP_K): + raise SystemExit(f"invalid top-k shape: {tuple(topk_ids.shape)}") + torch.testing.assert_close( + topk_weights.sum(dim=-1), + torch.ones(num_tokens, device=device), + atol=1e-5, + rtol=1e-5, + ) + torch.testing.assert_close(combined_weights, topk_weights) + torch.testing.assert_close(combined_ids, topk_ids) + additive_median = gate_time["median"] + topk_time["median"] + row = { + "num_tokens": num_tokens, + "gate_linear_time_ms": gate_time, + "routing_topk_time_ms": topk_time, + "gate_plus_topk_time_ms": combined_time, + "median_nonadditivity_ratio": ( + combined_time["median"] / additive_median + if additive_median > 0 + else 1.0 + ), + } + rows.append(row) + print(json.dumps(row, sort_keys=True), flush=True) + finally: + destroy_model_parallel() + destroy_distributed_environment() payload = { "schema_version": "qwen30_vllm020_router_raw.v1",