From 1fb777fc113001bffe39e26967881834ee122795 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Thu, 16 Jul 2026 21:48:52 +0800 Subject: [PATCH] Set vLLM config for collective profiling --- .../profile_vllm020_allreduce.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_allreduce.py b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_allreduce.py index 8665b9c..33fc9e8 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_allreduce.py +++ b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_allreduce.py @@ -69,6 +69,7 @@ def main() -> None: tensor_model_parallel_all_reduce, ) from vllm.distributed.parallel_state import get_tp_group + from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config rank = int(os.environ["RANK"]) local_rank = int(os.environ["LOCAL_RANK"]) @@ -79,7 +80,11 @@ def main() -> None: torch.accelerator.set_device_index(device) torch.set_default_device(device) init_distributed_environment() - initialize_model_parallel(tensor_model_parallel_size=world_size) + vllm_config = VllmConfig( + parallel_config=ParallelConfig(tensor_parallel_size=world_size) + ) + with set_current_vllm_config(vllm_config): + initialize_model_parallel(tensor_model_parallel_size=world_size) rows: list[dict[str, Any]] = [] expected_sum = world_size * (world_size + 1) / 2