Set vLLM config for collective profiling

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