gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Profile MoE TP-local shards without collectives

Browse Source
This commit is contained in:
Gahow Wang
2026-07-16 21:41:38 +08:00
parent 9bbfb87a85
commit 4fed4329cf
2 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
runs/frontier-qwen30-vllm020-profile-v1/jobs_moe_smoke.toml
View File

@@ -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"

10
runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_moe.py
View File

@@ -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,
}