eb24dc4a45 · tduncan/Nvidia-Optimization-Stack-VLLM

.buildkite/test-pipeline.yaml

+1

@@ -503,6 +503,7 @@ steps:

503 - entrypoints/llm/test_collective_rpc.py

504 commands:

505 - pytest -v -s entrypoints/llm/test_collective_rpc.py

506 + - VLLM_USE_V1=1 torchrun --nproc-per-node=2 distributed/test_torchrun_example.py

507 - torchrun --nproc-per-node=2 distributed/test_torchrun_example.py

508 - pytest -v -s ./compile/test_basic_correctness.py

509 - pytest -v -s ./compile/test_wrapper.py

tests/distributed/test_torchrun_example.py

+6

@@ -48,6 +48,12 @@ test_consistent_across_ranks(

48 test_consistent_across_ranks(

49 llm.llm_engine.vllm_config.cache_config.num_gpu_blocks)

50

51 +# make sure we can access the model parameters from the calling process

52 +# of the `LLM` instance.

53 +params = list(llm.llm_engine.model_executor.driver_worker.worker.model_runner.

54 + model.parameters())

55 +test_consistent_across_ranks(len(params))

56 +

57 # all ranks should have the same outputs

58 for output in outputs:

59 prompt = output.prompt

tests/v1/engine/test_engine_core.py

+4 -2

@@ -5,6 +5,7 @@ import threading

5 import time

6 import uuid

7 from concurrent.futures import Future

8 +from typing import List

9

10 import pytest

11 from transformers import AutoTokenizer

@@ -211,8 +212,9 @@ def test_engine_core_concurrent_batches(monkeypatch):

212

213 class DummyExecutor(UniProcExecutor):

214

214 - def initialize(self, kv_cache_config: KVCacheConfig) -> None:

215 - super().initialize(kv_cache_config)

215 + def initialize_from_config(

216 + self, kv_cache_configs: List[KVCacheConfig]) -> None:

217 + super().initialize_from_config(kv_cache_configs)

218

219 # This executor actually can only run 1 batch at a time

220 self.semaphore = threading.Semaphore(1)

vllm/config.py

+5

@@ -1407,6 +1407,11 @@ class ParallelConfig:

1407 self.data_parallel_master_port = envs.VLLM_DP_MASTER_PORT

1408 self.world_size_across_dp = self.world_size * self.data_parallel_size

1409

1410 + if self.distributed_executor_backend == "external_launcher":

1411 + import os

1412 + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"

1413 + logger.info("Disabling V1 multiprocessing for external launcher.")

1414 +

1415 ray_only_devices = ["tpu"]

1416 from vllm.platforms import current_platform

1417 if (current_platform.device_type in ray_only_devices

vllm/executor/ray_distributed_executor.py

+1 -1

@@ -541,7 +541,7 @@ class RayDistributedExecutor(DistributedExecutorBase):

541 # and the TP group executes in SPMD fashion.

542 if self.use_v1:

543 outputs = [

544 - worker.execute_model.

544 + worker.execute_model_ray.

545 bind( # type: ignore[attr-defined]

546 outputs[i]) for i, worker in enumerate(tp_group)

547 ]

vllm/executor/ray_utils.py

+3 -1

@@ -112,10 +112,12 @@ try:

112 torch.cuda.set_device(self.worker.device)

113 self.compiled_dag_cuda_device_set = True

114

115 - def execute_model(

115 + def execute_model_ray(

116 self,

117 scheduler_output: "SchedulerOutput",

118 ) -> "ModelRunnerOutput":

119 + # this method is used to compile ray CG,

120 + # and it needs a special logic of self.setup_device_if_necessary()

121 self.setup_device_if_necessary()

122 assert self.worker is not None, "Worker is not initialized"

123 if isinstance(scheduler_output, tuple):

vllm/executor/uniproc_executor.py

+4 -3

@@ -93,9 +93,10 @@ class ExecutorWithExternalLauncher(UniProcExecutor):

93 ("ExecutorWithExternalLauncher needs deterministic "

94 "execution, so it"

95 "does not support delay_factor in scheduling")

96 - assert not envs.VLLM_USE_V1, \

97 - ("V1 architecture cannot guarantee deterministic execution, "

98 - "so it is not supported in ExecutorWithExternalLauncher.")

96 + if envs.VLLM_USE_V1:

97 + assert not envs.VLLM_ENABLE_V1_MULTIPROCESSING, \

98 + ("To get deterministic execution in V1, "

99 + "please set VLLM_ENABLE_V1_MULTIPROCESSING=0")

100 self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config,

101 rpc_rank=0)

102 # engines are launched in torchrun-compatible launchers

vllm/v1/engine/core.py

+1 -1

@@ -110,7 +110,7 @@ class EngineCore:

110 num_cpu_blocks = 0

111

112 # Initialize kv cache and warmup the execution

113 - self.model_executor.initialize(kv_cache_configs)

113 + self.model_executor.initialize_from_config(kv_cache_configs)

114

115 elapsed = time.time() - start

116 logger.info(("init engine (profile, create kv cache, "

vllm/v1/engine/llm_engine.py

+7 -2

@@ -4,10 +4,10 @@ from typing import Dict, List, Mapping, Optional, Type, Union

4

5 from typing_extensions import TypeVar

6

7 +import vllm.envs as envs

8 from vllm.config import ParallelConfig, VllmConfig

9 from vllm.engine.arg_utils import EngineArgs

10 from vllm.engine.metrics_types import StatLoggerBase

10 -from vllm.envs import VLLM_ENABLE_V1_MULTIPROCESSING

11 from vllm.inputs import INPUT_REGISTRY, InputRegistry, PromptType

12 from vllm.logger import init_logger

13 from vllm.lora.request import LoRARequest

@@ -44,6 +44,7 @@ class LLMEngine:

44 use_cached_outputs: bool = False,

45 multiprocess_mode: bool = False,

46 ) -> None:

47 + self.vllm_config = vllm_config

48 self.model_config = vllm_config.model_config

49 self.cache_config = vllm_config.cache_config

50

@@ -83,6 +84,10 @@ class LLMEngine:

84 log_stats=False, # FIXME: implement

85 )

86

87 + if not multiprocess_mode:

88 + # for v0 compatibility

89 + self.model_executor = self.engine_core.engine_core.model_executor # type: ignore

90 +

91 @classmethod

92 def from_engine_args(

93 cls,

@@ -97,7 +102,7 @@ class LLMEngine:

102 vllm_config = engine_args.create_engine_config(usage_context)

103 executor_class = Executor.get_class(vllm_config)

104

100 - if VLLM_ENABLE_V1_MULTIPROCESSING:

105 + if envs.VLLM_ENABLE_V1_MULTIPROCESSING:

106 logger.debug("Enabling multiprocessing for LLMEngine.")

107 enable_multiprocessing = True

108

vllm/v1/executor/abstract.py

+17 -3

@@ -3,6 +3,9 @@

3 from concurrent.futures import Future

4 from typing import List, Type, Union

5

6 +import torch

7 +import torch.distributed as dist

8 +

9 from vllm.config import VllmConfig

10 from vllm.executor.executor_base import ExecutorBase

11 from vllm.executor.uniproc_executor import ( # noqa

@@ -49,12 +52,14 @@ class Executor(ExecutorBase):

52 f"{distributed_executor_backend}")

53 return executor_class

54

52 - def initialize(self, kv_cache_configs: List[KVCacheConfig]) -> None:

55 + def initialize_from_config(self,

56 + kv_cache_configs: List[KVCacheConfig]) -> None:

57 """

58 Initialize the KV caches and begin the model execution loop of the

59 underlying workers.

60 """

57 - self.collective_rpc("initialize_cache", args=(kv_cache_configs, ))

61 + self.collective_rpc("initialize_from_config",

62 + args=(kv_cache_configs, ))

63 self.collective_rpc("compile_or_warm_up_model")

64

65 def determine_available_memory(self) -> int: # in bytes

@@ -89,4 +94,13 @@ class UniProcExecutor(UniProcExecutorV0, Executor):

94

95

96 class ExecutorWithExternalLauncher(ExecutorWithExternalLauncherV0, Executor):

92 - pass

97 +

98 + def determine_available_memory(self) -> int: # in bytes

99 + # same as determine_num_available_blocks in v0,

100 + # we need to get the min across all ranks.

101 + memory = super().determine_available_memory()

102 + from vllm.distributed.parallel_state import get_world_group

103 + cpu_group = get_world_group().cpu_group

104 + memory_tensor = torch.tensor([memory], device="cpu", dtype=torch.int64)

105 + dist.all_reduce(memory_tensor, group=cpu_group, op=dist.ReduceOp.MIN)

106 + return memory_tensor.item()

vllm/v1/executor/multiproc_executor.py

+3 -2

@@ -216,9 +216,10 @@ class WorkerProc:

216 "local_rank": local_rank,

217 "rank": rank,

218 "distributed_init_method": distributed_init_method,

219 + "is_driver_worker": rank == 0,

220 }

221 wrapper.init_worker(all_kwargs)

221 - self.worker = wrapper.worker

222 + self.worker = wrapper

223

224 pid = os.getpid()

225 _add_prefix(sys.stdout, f"VllmWorker rank={rank}", pid)

@@ -239,7 +240,7 @@ class WorkerProc:

240 ready_socket.send_string(WorkerProc.READY_STR)

241 ready_socket.send(payload)

242

242 - wrapper.init_device()

243 + self.worker.init_device()

244 self.worker.load_model()

245

246 @staticmethod

vllm/v1/worker/gpu_worker.py

+3 -4

@@ -2,7 +2,7 @@

2 """A GPU worker class."""

3 import gc

4 import os

5 -from typing import TYPE_CHECKING, List, Optional

5 +from typing import TYPE_CHECKING, Optional

6

7 import torch

8 import torch.distributed

@@ -185,9 +185,8 @@ class Worker(WorkerBase):

185 def get_kv_cache_spec(self) -> KVCacheSpec:

186 return self.model_runner.get_kv_cache_spec()

187

188 - def initialize_cache(self, kv_cache_configs: List[KVCacheConfig]) -> None:

188 + def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:

189 """Allocate GPU KV cache with the specified kv_cache_config."""

190 - kv_cache_config = kv_cache_configs[self.rank]

190 if self.vllm_config.model_config.enable_sleep_mode:

191 allocator = CuMemAllocator.get_instance()

192 context = allocator.use_memory_pool(tag="kv_cache")

@@ -225,7 +224,7 @@ class Worker(WorkerBase):

224 scheduler_output: "SchedulerOutput",

225 ) -> Optional[ModelRunnerOutput]:

226 output = self.model_runner.execute_model(scheduler_output)

228 - return output if self.rank == 0 else None

227 + return output if self.is_driver_worker else None

228

229 def profile(self, is_start: bool = True):

230 if self.profiler is None:

vllm/v1/worker/tpu_worker.py

+3 -3

@@ -36,6 +36,7 @@ class TPUWorker:

36 distributed_init_method: str,

37 is_driver_worker: bool = False,

38 ):

39 + self.is_driver_worker = is_driver_worker

40 self.vllm_config = vllm_config

41 self.model_config = vllm_config.model_config

42 self.cache_config = vllm_config.cache_config

@@ -151,7 +152,7 @@ class TPUWorker:

152 scheduler_output: "SchedulerOutput",

153 ) -> Optional[ModelRunnerOutput]:

154 output = self.model_runner.execute_model(scheduler_output)

154 - return output if self.rank == 0 else None

155 + return output if self.is_driver_worker else None

156

157 def load_model(self) -> None:

158 self.model_runner.load_model()

@@ -170,9 +171,8 @@ class TPUWorker:

171 def get_kv_cache_spec(self) -> KVCacheSpec:

172 return self.model_runner.get_kv_cache_spec()

173

173 - def initialize_cache(self, kv_cache_configs: List[KVCacheConfig]) -> None:

174 + def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:

175 """Allocate GPU KV cache with the specified kv_cache_config."""

175 - kv_cache_config = kv_cache_configs[self.rank]

176 self.model_runner.initialize_kv_cache(kv_cache_config)

177

178 def check_health(self) -> None:

vllm/worker/worker_base.py

+9 -2

@@ -567,6 +567,10 @@ class WorkerWrapperBase:

567 self.worker = worker_class(**kwargs)

568 assert self.worker is not None

569

570 + def initialize_from_config(self, kv_cache_configs: List[Any]) -> None:

571 + kv_cache_config = kv_cache_configs[self.rpc_rank]

572 + self.worker.initialize_from_config(kv_cache_config) # type: ignore

573 +

574 def init_device(self):

575 with set_current_vllm_config(self.vllm_config):

576 # To make vLLM config available during device initialization

@@ -574,8 +578,11 @@ class WorkerWrapperBase:

578

579 def execute_method(self, method: Union[str, bytes], *args, **kwargs):

580 try:

577 - target = self if self.worker is None else self.worker

578 - return run_method(target, method, args, kwargs)

581 + # method resolution order:

582 + # if a method is defined in this class, it will be called directly.

583 + # otherwise, since we define `__getattr__` and redirect attribute

584 + # query to `self.worker`, the method will be called on the worker.

585 + return run_method(self, method, args, kwargs)

586 except Exception as e:

587 # if the driver worker also execute methods,

588 # exceptions in the rest worker may cause deadlock in rpc like ray

[v1] torchrun compatibility (#13642)