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