A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 16:00 UTC
[Misc] Use envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE (#15831)
Signed-off-by: Rui Qiao <ruisearch42@gmail.com> Co-authored-by: Cody Yu <hao.yu.cody@gmail.com> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk>
Rui Qiao committed
Apr 1, 2025 at 06:07 UTC
8dd41d6bcc6680b69819cd6518c44ff3242b73f1
2 files changed
+30
-17
vllm/envs.py
+13
-7
@@ -48,7 +48,7 @@ if TYPE_CHECKING:
48
VLLM_FUSED_MOE_CHUNK_SIZE: int = 64 * 1024
49
VLLM_USE_RAY_SPMD_WORKER: bool = False
50
VLLM_USE_RAY_COMPILED_DAG: bool = False
51
- VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL: bool = True
51
+ VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: str = "auto"
52
VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False
53
VLLM_WORKER_MULTIPROC_METHOD: str = "fork"
54
VLLM_ASSETS_CACHE: str = os.path.join(VLLM_CACHE_ROOT, "assets")
@@ -380,15 +380,21 @@ environment_variables: dict[str, Callable[[], Any]] = {
380
# (previously known as ADAG) API which optimizes the
381
# control plane overhead.
382
# Run vLLM with VLLM_USE_RAY_COMPILED_DAG=1 to enable it.
383
+ # Note that this variable is set to 1 in V1 by default
384
+ # when ray distributed executor is used.
385
"VLLM_USE_RAY_COMPILED_DAG":
386
lambda: bool(int(os.getenv("VLLM_USE_RAY_COMPILED_DAG", "0"))),
387
386
- # If the env var is set, it uses NCCL for communication in
387
- # Ray's Compiled Graph. This flag is ignored if
388
- # VLLM_USE_RAY_COMPILED_DAG is not set.
389
- "VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL":
390
- lambda: bool(int(os.getenv("VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL", "1"))
391
- ),
388
+ # If the env var is set, Ray Compiled Graph uses the specified
389
+ # channel type to communicate between workers belonging to
390
+ # different pipeline-parallel stages.
391
+ # Available options:
392
+ # - "auto": use the default channel type
393
+ # - "nccl": use NCCL for communication
394
+ # - "shm": use shared memory and gRPC for communication
395
+ # This flag is ignored if VLLM_USE_RAY_COMPILED_DAG is not set.
396
+ "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE":
397
+ lambda: os.getenv("VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", "auto"),
398
399
# If the env var is set, it enables GPU communication overlap
400
# (experimental feature) in Ray's Compiled Graph. This flag is ignored if
vllm/executor/ray_distributed_executor.py
+17
-10
@@ -79,7 +79,7 @@ class RayDistributedExecutor(DistributedExecutorBase):
79
80
# For TPU, avoid compiling NVIDIA's NCCL
81
if current_platform.is_tpu():
82
- os.environ["VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL"] = "0"
82
+ os.environ["VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE"] = "shm"
83
84
# If the env var is set, it uses the Ray's compiled DAG API
85
# which optimizes the control plane overhead.
@@ -546,10 +546,11 @@ class RayDistributedExecutor(DistributedExecutorBase):
546
"Run `pip install ray[cgraph]` to install it.")
547
548
cupy_spec = importlib.util.find_spec("cupy")
549
- if cupy_spec is None and envs.VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL:
549
+ if (cupy_spec is None
550
+ and envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE == "nccl"):
551
raise ValueError(
552
"cupy is not installed but required since "
552
- "VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL is set. "
553
+ "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE is set to 'nccl'. "
554
"Run `pip install ray[cgraph]` and check cupy installation.")
555
556
def _compiled_ray_dag(self, enable_asyncio: bool):
@@ -557,10 +558,17 @@ class RayDistributedExecutor(DistributedExecutorBase):
558
self._check_ray_cgraph_installation()
559
from ray.dag import InputNode, MultiOutputNode
560
560
- logger.info("VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL = %s",
561
- envs.VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL)
561
+ logger.info("VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE = %s",
562
+ envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE)
563
logger.info("VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM = %s",
564
envs.VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM)
565
+
566
+ channel_type = envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE
567
+ if channel_type not in ("auto", "nccl", "shm"):
568
+ raise ValueError(
569
+ "Invalid value for VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: "
570
+ f"{channel_type}. Valid values are: 'auto', 'nccl', or 'shm'.")
571
+
572
# Enlarge the default value of "RAY_CGRAPH_get_timeout" to 300 seconds
573
# (it is 10 seconds by default). This is a Ray environment variable to
574
# control the timeout of getting result from a compiled graph execution,
@@ -605,13 +613,12 @@ class RayDistributedExecutor(DistributedExecutorBase):
613
]
614
615
last_pp_rank = len(self.pp_tp_workers) - 1
608
- if pp_rank < last_pp_rank:
616
+ if (pp_rank < last_pp_rank and
617
+ envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE != "shm"):
618
# Specify how intermediate tensors should be passed
619
# between pp stages, no need to specify for the last
611
- # pp stage.
612
- transport = "nccl" \
613
- if envs.VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL \
614
- else "auto"
620
+ # pp stage or when using shared memory (the default).
621
+ transport = envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE
622
outputs = [
623
output.with_tensor_transport(transport=transport)
624
for output in outputs