@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 140 lines 5.15 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import queue
5 import random
6 import typing
7
8 import pytest
9 import torch
10 import torch.distributed as dist
11 import torch.multiprocessing as mp
12
13 import vllm.envs as envs
14 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
15 from vllm.distributed import cleanup_dist_env_and_memory
16 from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
17 from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
18 from vllm.distributed.parallel_state import (
19 get_tp_group,
20 init_distributed_environment,
21 initialize_model_parallel,
22 )
23 from vllm.engine.arg_utils import EngineArgs
24 from vllm.engine.llm_engine import LLMEngine
25 from vllm.platforms import current_platform
26 from vllm.utils.system_utils import update_environment_variables
27
28 torch.manual_seed(42)
29 random.seed(44)
30
31 test_size_elements = 1024 * 1024
32
33
34 def symm_mem_allreduce_worker(local_rank: int, world_size: int, q: mp.Queue):
35 monkeypatch = pytest.MonkeyPatch()
36 config = VllmConfig(parallel_config=ParallelConfig(tensor_parallel_size=world_size))
37
38 with monkeypatch.context() as m, set_current_vllm_config(config):
39 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
40 dtype = torch.bfloat16
41 device = torch.device(f"cuda:{local_rank}")
42 torch.accelerator.set_device_index(device)
43 torch.set_default_device(device)
44 torch.set_default_dtype(dtype)
45 update_environment_variables(
46 {
47 "RANK": str(local_rank),
48 "LOCAL_RANK": str(local_rank),
49 "WORLD_SIZE": str(world_size),
50 "MASTER_ADDR": "localhost",
51 "MASTER_PORT": "12345",
52 }
53 )
54
55 init_distributed_environment()
56 initialize_model_parallel(tensor_model_parallel_size=world_size)
57
58 cuda_communicator = typing.cast(
59 CudaCommunicator, get_tp_group().device_communicator
60 )
61 symm_mem_comm = cuda_communicator.symm_mem_comm
62 if symm_mem_comm is None or symm_mem_comm.disabled:
63 # can't use skip under multiprocessing
64 q.put("SymmMemCommunicator is not available or disabled.")
65 return
66
67 inp_direct_symm_mem = torch.randint(
68 1, 23, (test_size_elements,), dtype=dtype, device=device
69 )
70 if not symm_mem_comm.should_use_symm_mem(inp_direct_symm_mem):
71 # can't use skip under multiprocessing
72 q.put("SymmMemCommunicator isn't used for this world and input size.")
73 return
74
75 original_inp_direct_symm_mem = inp_direct_symm_mem.clone()
76 out_direct_symm_mem = symm_mem_comm.all_reduce(inp_direct_symm_mem)
77 assert out_direct_symm_mem is not None
78
79 group = get_tp_group().device_group
80 dist.all_reduce(original_inp_direct_symm_mem, group=group)
81 torch.testing.assert_close(
82 out_direct_symm_mem, original_inp_direct_symm_mem, atol=2.5, rtol=0.1
83 )
84
85 # Test tensor_model_parallel_all_reduce which should use symm_mem
86 inp_tensor_parallel = torch.randint(
87 -23, 1, (test_size_elements,), dtype=dtype, device=device
88 )
89 original_inp_tensor_parallel = inp_tensor_parallel.clone()
90 out_tensor_parallel = tensor_model_parallel_all_reduce(inp_tensor_parallel)
91 dist.all_reduce(original_inp_tensor_parallel, group=group)
92 torch.testing.assert_close(
93 out_tensor_parallel, original_inp_tensor_parallel, atol=2.5, rtol=0.1
94 )
95
96
97 @pytest.mark.skipif(
98 not current_platform.is_cuda(),
99 reason="SymmMemAllreduce is only available for CUDA platforms.",
100 )
101 @pytest.mark.parametrize("tp_size", [2])
102 @pytest.mark.parametrize("pipeline_parallel_size", [1])
103 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
104 def test_symm_mem_allreduce(
105 monkeypatch: pytest.MonkeyPatch, tp_size, pipeline_parallel_size
106 ):
107 world_size = tp_size * pipeline_parallel_size
108 if world_size > torch.accelerator.device_count():
109 pytest.skip("Not enough GPUs to run the test.")
110 q = mp.get_context("spawn").Queue()
111 mp.spawn(symm_mem_allreduce_worker, args=(world_size, q), nprocs=world_size)
112 try:
113 val = q.get(timeout=1)
114 except queue.Empty:
115 val = None
116 finally:
117 cleanup_dist_env_and_memory()
118 if val is not None:
119 pytest.skip(val)
120
121
122 @pytest.mark.skipif(
123 not current_platform.is_cuda(),
124 reason="SymmMemAllreduce is only available for CUDA platforms.",
125 )
126 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
127 def test_dp_with_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch):
128 world_size = 4
129 if world_size > torch.accelerator.device_count():
130 pytest.skip("Not enough GPUs to run the test.")
131 # Verify that the DataParallel runs without error
132 engine_args = EngineArgs(
133 model="distilbert/distilgpt2",
134 enforce_eager=True,
135 enable_prefix_caching=True,
136 data_parallel_size=2,
137 tensor_parallel_size=2,
138 data_parallel_backend="mp",
139 )
140 LLMEngine.from_engine_args(engine_args)