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

Set vLLM config for collective profiling

Browse Source
This commit is contained in:
Gahow Wang
2026-07-16 21:48:52 +08:00
parent 4fed4329cf
commit 1fb777fc11
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_allreduce.py
View File

@@ -69,6 +69,7 @@ def main() -> None:
tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce,
) )
from vllm.distributed.parallel_state import get_tp_group from vllm.distributed.parallel_state import get_tp_group
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
rank = int(os.environ["RANK"]) rank = int(os.environ["RANK"])
local_rank = int(os.environ["LOCAL_RANK"]) local_rank = int(os.environ["LOCAL_RANK"])
@@ -79,6 +80,10 @@ def main() -> None:
torch.accelerator.set_device_index(device) torch.accelerator.set_device_index(device)
torch.set_default_device(device) torch.set_default_device(device)
init_distributed_environment() init_distributed_environment()
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) initialize_model_parallel(tensor_model_parallel_size=world_size)
rows: list[dict[str, Any]] = [] rows: list[dict[str, Any]] = []