From 4fed4329cfe9383e294719886600cbaf5f3a1212 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Thu, 16 Jul 2026 21:41:38 +0800 Subject: [PATCH] Profile MoE TP-local shards without collectives --- .../jobs_moe_smoke.toml | 6 +++--- .../profile_vllm020_moe.py | 10 +++++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/runs/frontier-qwen30-vllm020-profile-v1/jobs_moe_smoke.toml b/runs/frontier-qwen30-vllm020-profile-v1/jobs_moe_smoke.toml index 48cd8db..f8403b3 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/jobs_moe_smoke.toml +++ b/runs/frontier-qwen30-vllm020-profile-v1/jobs_moe_smoke.toml @@ -1,18 +1,18 @@ version = 1 [[jobs]] -name = "qwen30-vllm020-moe-smoke-20260716-v1" +name = "qwen30-vllm020-moe-smoke-20260716-v2-local-shard" 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_moe_smoke.sh" -artifacts = ["artifacts/moe-smoke"] +artifacts = ["artifacts/moe-smoke-v2"] [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/moe-smoke" +OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/moe-smoke-v2" 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_moe.py b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_moe.py index 927b316..74da2d1 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_moe.py +++ b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_moe.py @@ -199,8 +199,12 @@ def main() -> None: in_dtype=torch.bfloat16, max_num_tokens=max_num_tokens, ) + # This process profiles one TP-local weight shard. Keep the global + # runtime context single-rank so vLLM does not initialize a collective; + # the action-conditioned shard size remains explicit in moe_config and + # the real TP2/TP4 all-reduce is profiled in a separate multi-GPU run. vllm_config = VllmConfig( - parallel_config=ParallelConfig(tensor_parallel_size=tp) + parallel_config=ParallelConfig(tensor_parallel_size=1) ) with set_current_vllm_config(vllm_config): backend, experts_cls = select_unquantized_moe_backend(moe_config) @@ -348,8 +352,8 @@ def main() -> None: "norm_topk_prob": True, }, "measurement_scope": ( - "vLLM modular MoE prepare+FlashInfer CUTLASS experts+finalize; " - "router linear/top-k and TP all-reduce excluded" + "one TP-local weight shard: vLLM modular MoE prepare+FlashInfer " + "CUTLASS experts+finalize; router linear/top-k and TP all-reduce excluded" ), "rows": rows, }