diff --git a/docs/configuration/conserving_memory.md b/docs/configuration/conserving_memory.md index 4d5c961af98fd..2c073603f68ab 100644 --- a/docs/configuration/conserving_memory.md +++ b/docs/configuration/conserving_memory.md @@ -19,7 +19,20 @@ llm = LLM(model="ibm-granite/granite-3.1-8b-instruct", To ensure that vLLM initializes CUDA correctly, you should avoid calling related functions (e.g. [torch.cuda.set_device][]) before initializing vLLM. Otherwise, you may run into an error like `RuntimeError: Cannot re-initialize CUDA in forked subprocess`. - To control which devices are used, please instead set the `CUDA_VISIBLE_DEVICES` environment variable. + + To control which devices are used, you can either set the `CUDA_VISIBLE_DEVICES` + environment variable, pass the `gpu_ids` parameter to the [LLM] constructor, + or use the `--gpu-ids` option with `vllm serve`. + + ```python + from vllm import LLM + + # Use GPUs 0 and 2 for execution without setting CUDA_VISIBLE_DEVICES env var + llm = LLM( + model="your-model", + gpu_ids=[0, 2], + ) + ``` !!! note With tensor parallelism enabled, each process will read the whole model and split it into chunks, which makes the disk reading time even longer (proportional to the size of tensor parallelism). diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index d25105cbb789f..eaa96f9bb251f 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -38,6 +38,9 @@ class ServeSubcommand(CLISubcommand): @staticmethod def cmd(args: argparse.Namespace) -> None: + # Allow overriding visible GPUs via --gpu-ids (comma-separated or single int) + if hasattr(args, 'gpu_ids') and args.gpu_ids is not None: + os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_ids # If model is specified in CLI (as positional arg), it takes precedence if hasattr(args, 'model_tag') and args.model_tag is not None: args.model = args.model_tag @@ -98,6 +101,12 @@ class ServeSubcommand(CLISubcommand): help="Read CLI options from a config file. " "Must be a YAML with the following options: " "https://docs.vllm.ai/en/latest/configuration/serve_args.html") + serve_parser.add_argument( + "--gpu-ids", + type=str, + default=None, + help="Comma-separated GPU IDs or a single GPU ID to use for vLLM serve. " + "Overrides CUDA_VISIBLE_DEVICES.") serve_parser = make_arg_parser(serve_parser) show_filtered_argument_or_group_from_help(serve_parser, ["serve"]) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 16c051d61de32..a0084ede3e61d 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -9,6 +9,7 @@ from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Optional, Union, cast, overload) import cloudpickle +import os import torch.nn as nn from pydantic import ValidationError from tqdm.auto import tqdm @@ -75,6 +76,9 @@ class LLM: skip_tokenizer_init: If true, skip initialization of tokenizer and detokenizer. Expect valid prompt_token_ids and None for prompt from the input. + gpu_ids: A list of GPU device IDs or a comma-separated string to use + for vLLM execution. Overrides the CUDA_VISIBLE_DEVICES environment + variable. trust_remote_code: Trust remote code (e.g., from HuggingFace) when downloading the model and tokenizer. allowed_local_media_path: Allowing API requests to read local images @@ -170,6 +174,7 @@ class LLM: tokenizer: Optional[str] = None, tokenizer_mode: TokenizerMode = "auto", skip_tokenizer_init: bool = False, + gpu_ids: Optional[Union[Sequence[int], str]] = None, trust_remote_code: bool = False, allowed_local_media_path: str = "", tensor_parallel_size: int = 1, @@ -198,6 +203,13 @@ class LLM: if "disable_log_stats" not in kwargs: kwargs["disable_log_stats"] = True + # Allow specifying GPU device IDs without using CUDA_VISIBLE_DEVICES env var + if gpu_ids is not None: + # gpu_ids can be a sequence of ints or a string + os.environ["CUDA_VISIBLE_DEVICES"] = ( + ",".join(map(str, gpu_ids)) if isinstance(gpu_ids, (list, tuple)) + else str(gpu_ids) + ) if "worker_cls" in kwargs: worker_cls = kwargs["worker_cls"] # if the worker_cls is not qualified string name,