Working with LLMs#

Use the ray.data.llm module to run scalable batch inference on Ray Data datasets. The module can run a vLLM or SGLang inference engine directly, query OpenAI-compatible endpoints through HttpRequestProcessorConfig, or send requests to a Ray Serve deployment through ServeDeploymentProcessorConfig.

To get started, see the following sections:

Quickstart: Run your first batch inference job.
Architecture: Understand the processor pipeline.
Scaling: Scale your LLM stage to multiple replicas.

For common use cases, see the following sections:

Text generation: Generate chat completions with LLMs.
Embeddings: Generate text embeddings.
Classification: Run content classifiers and sentiment analyzers.
Multimodality: Run batch inference with vision language models (VLMs) and omni-modal models on multimodal data.
OpenAI-compatible endpoints: Query deployed models.
Serve deployments: Share vLLM engines across processors.
Custom tokenizers: Use vLLM tokenizers for models that Hugging Face doesn’t support.

To troubleshoot and tune your pipeline, see the following sections:

Troubleshooting: Resolve GPU memory and model loading issues.
Advanced configuration: Configure parallelism, per-stage tuning, LoRA, and batch concurrency.

Quickstart: Run batch inference with vLLM#

This example shows the minimal setup to run vLLM batch inference on a dataset.

Note

This quickstart requires a GPU because vLLM is GPU-accelerated.

First, install Ray Data with LLM support:

pip install -U "ray[data, llm]>=2.53.0"

The following minimal example runs batch inference:

import ray
from ray.data.llm import vLLMEngineProcessorConfig, build_processor

# Initialize Ray
ray.init()

# simple dataset
ds = ray.data.from_items([
    {"prompt": "What is machine learning?"},
    {"prompt": "Explain neural networks in one sentence."},
])

# Minimal vLLM configuration
config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    concurrency=1,  # 1 vLLM engine replica
    batch_size=32,  # 32 samples per batch
    engine_kwargs={
        "max_model_len": 4096, # Fit into test GPU memory
    }
)

# Build processor
# preprocess: converts input row to format expected by vLLM (OpenAI chat format)
# postprocess: extracts generated text from vLLM output
processor = build_processor(
    config,
    preprocess=lambda row: {
        "messages": [{"role": "user", "content": row["prompt"]}],
        "sampling_params": {"temperature": 0.7, "max_tokens": 100},
    },
    postprocess=lambda row: {
        "prompt": row["prompt"],
        "response": row["generated_text"],
    },
)

# inference
ds = processor(ds)

# iterate through the results
for result in ds.iter_rows():
    print(f"Q: {result['prompt']}")
    print(f"A: {result['response']}\n")

# Alternative ways to get results:
# results = ds.take(10)  # Get first 10 results
# ds.show(limit=5)       # Print first 5 results
# ds.write_parquet("output.parquet")  # Save to file

This example does the following:

Creates a simple dataset with prompts.
Configures a vLLM processor with minimal settings.
Builds a processor that converts prompts to OpenAI chat format during preprocessing and extracts the generated text during postprocessing.
Runs inference on the dataset.
Iterates through the results.

The processor expects input rows with a prompt field and outputs rows with both prompt and response fields. Consume the results with iter_rows(), take(), or show(), or save them to files with write_parquet().

For more configuration options and features, see the following sections.

How does the processor pipeline work?#

Ray Data LLM transforms your data through a multi-stage processor pipeline that runs LLM inference. Understanding this architecture helps you optimize performance and debug issues.

Input Dataset
     |
     v
- Preprocess (Custom Function)
- PrepareMultimodal (Optional, for VLM / Omni models)
- ChatTemplate (Applies chat template to messages)
- Tokenize (Optional -- converts text to token IDs)
- LLM Engine (vLLM/SGLang inference on GPU)
- Detokenize (Optional -- converts token IDs back to text)
- Postprocess (Custom Function)
     |
     v
Output Dataset

Each stage does the following:

Preprocess: Your custom function that transforms input rows into the format that downstream stages expect, typically OpenAI chat format with messages.
PrepareMultimodal: Extracts and prepares multimodal inputs. Enable it with prepare_multimodal_stage=True.
ChatTemplate: Applies the model’s chat template to convert messages into a prompt string.
Tokenize: Converts the prompt string into token IDs for the model.
LLM Engine: The inference stage that runs vLLM or SGLang on GPU or TPU accelerators.
Detokenize: Converts output token IDs back to readable text.
Postprocess: Your custom function that extracts and formats the final output.

Each stage runs as a separate Ray actor pool, so you can scale and allocate resources for each stage independently. By default, all CPU and GPU stages use autoscaling actor pools, except the ServeDeployment stage, which uses a fixed pool.

Scale to multiple GPUs#

To scale the LLM stage horizontally to multiple GPU replicas, set the concurrency parameter:

config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={
        "enable_chunked_prefill": True,
        "max_num_batched_tokens": 4096,
        "max_model_len": 16384,
    },
    concurrency=10,
    batch_size=64,
)

Each replica runs an independent inference engine. Set concurrency to match the number of available GPUs or GPU nodes.

By default, when you set concurrency to an integer n, GPU stages autoscale from 1 to n actors. To use a fixed pool of n actors, set concurrency to (n, n).

config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={
        "enable_chunked_prefill": True,
        "max_num_batched_tokens": 4096,
        "max_model_len": 16384,
    },
    concurrency=(10, 10),
    batch_size=64,
)

Generate text#

Use vLLMEngineProcessorConfig or SGLangEngineProcessorConfig for chat completions and text generation tasks.

The key configuration options include the following:

model_source: The Hugging Face model ID or the path to the model weights.
concurrency: The number of vLLM engine replicas, typically one per GPU node.
batch_size: The number of rows per batch. Reduce it if you hit memory limits.

# Basic vLLM configuration
config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={
        "enable_chunked_prefill": True,
        "max_num_batched_tokens": 4096,  # Reduce if CUDA OOM occurs
        "max_model_len": 4096,  # Constrain to fit test GPU memory
    },
    concurrency=1,
    batch_size=64,
)

For gated models that require authentication, pass your Hugging Face token through runtime_env:

# Configuration with Hugging Face token
config_with_token = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    runtime_env={"env_vars": {"HF_TOKEN": "your_huggingface_token"}},
    concurrency=1,
    batch_size=64,
)

Run batch inference on multimodal data#

Ray Data LLM also supports batch inference with vision language models and omni-modal models on multimodal data. To run multimodal batch inference, make the following two changes to the previous example:

Set prepare_multimodal_stage=True in the vLLMEngineProcessorConfig.
Prepare multimodal data inside the preprocessor.

Run image batch inference with a vision language model#

First, load a vision dataset:

    """
    Load vision dataset from Hugging Face.

    This function loads the LMMs-Eval-Lite dataset which contains:
    - Images with associated questions
    - Multiple choice answers
    - Various visual reasoning tasks
    """
    try:
        from huggingface_hub import HfFileSystem

        # Load "LMMs-Eval-Lite" dataset from Hugging Face using HfFileSystem
        path = "hf://datasets/lmms-lab/LMMs-Eval-Lite/coco2017_cap_val/"
        fs = HfFileSystem()
        vision_dataset = ray.data.read_parquet(path, filesystem=fs)

        return vision_dataset
    except ImportError:
        print(
            "huggingface_hub package not available. Install with: pip install huggingface_hub"
        )
        return None
    except Exception as e:
        print(f"Error loading dataset: {e}")
        return None

Next, configure the VLM processor with the essential settings:

vision_processor_config = vLLMEngineProcessorConfig(
    model_source="Qwen/Qwen2.5-VL-3B-Instruct",
    engine_kwargs=dict(
        tensor_parallel_size=1,
        pipeline_parallel_size=1,
        max_model_len=4096,
        trust_remote_code=True,
        limit_mm_per_prompt={"image": 1},
    ),
    batch_size=16,
    concurrency=1,
    prepare_multimodal_stage=True,
)

Define preprocessing and postprocessing functions that convert dataset rows into the format the VLM expects and extract the model responses. In the preprocessor, structure image data as part of an OpenAI-compatible message. You can pass either an image URL or a PIL.Image.Image object.

"""Supported image input formats: image URL, PIL Image object"""
{
    "messages": [
        {
            "role": "system",
            "content": "Provide a detailed description of the image."
        },
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "Describe what happens in this image."},
                # Option 1: Provide image URL
                {"type": "image_url", "image_url": {"url": "https://example.com/image.jpg"}},
                # Option 2: Provide PIL Image object
                {"type": "image_pil", "image_pil": PIL.Image.open("path/to/image.jpg")}
            ]
        },
    ]
}

def vision_preprocess(row: dict) -> dict:
    """
    Preprocessing function for vision-language model inputs.

    Converts dataset rows into the format expected by the VLM:
    - System prompt for analysis instructions
    - User message with text and image content
    - Multiple choice formatting
    - Sampling parameters
    """
    choice_indices = ["A", "B", "C", "D", "E", "F", "G", "H"]

    return {
        "messages": [
            {
                "role": "system",
                "content": (
                    "Analyze the image and question carefully, using step-by-step reasoning. "
                    "First, describe any image provided in detail. Then, present your reasoning. "
                    "And finally your final answer in this format: Final Answer: <answer> "
                    "where <answer> is: The single correct letter choice A, B, C, D, E, F, etc. when options are provided. "
                    "Only include the letter. Your direct answer if no options are given, as a single phrase or number. "
                    "IMPORTANT: Remember, to end your answer with Final Answer: <answer>."
                ),
            },
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": row["question"] + "\n\n"},
                    {
                        "type": "image_pil",
                        "image_pil": Image.open(BytesIO(row["image"]["bytes"])),
                    },
                    {
                        "type": "text",
                        "text": "\n\nChoices:\n"
                        + "\n".join(
                            [
                                f"{choice_indices[i]}. {choice}"
                                for i, choice in enumerate(row["answer"])
                            ]
                        ),
                    },
                ],
            },
        ],
        "sampling_params": {
            "temperature": 0.3,
            "max_tokens": 150,
            "detokenize": False,
        },
        # Include original data for reference
        "original_data": {
            "question": row["question"],
            "answer_choices": row["answer"],
            "image_size": row["image"].get("width", 0) if row["image"] else 0,
        },
    }


def vision_postprocess(row: dict) -> dict:
    return {
        "resp": row["generated_text"],
    }

Finally, run the VLM inference:

    """Run the complete VLM example workflow."""
    config = create_vlm_config()
    vision_dataset = load_vision_dataset()

    if vision_dataset:
        # Build processor with preprocessing and postprocessing
        processor = build_processor(
            config, preprocess=vision_preprocess, postprocess=vision_postprocess
        )

        print("VLM processor configured successfully")
        print(f"Model: {config.model_source}")
        result = processor(vision_dataset).take_all()
        return config, processor, result

Run video batch inference with a vision language model#

First, load a video dataset:

    """
    Load video dataset from ShareGPTVideo Hugging Face dataset.
    """
    try:
        from huggingface_hub import hf_hub_download
        import tarfile
        from pathlib import Path

        dataset_name = "ShareGPTVideo/train_raw_video"

        tar_path = hf_hub_download(
            repo_id=dataset_name,
            filename="activitynet/chunk_0.tar.gz",
            repo_type="dataset",
        )

        extract_dir = "/tmp/sharegpt_videos"
        os.makedirs(extract_dir, exist_ok=True)

        if not any(Path(extract_dir).glob("*.mp4")):
            with tarfile.open(tar_path, "r:gz") as tar:
                tar.extractall(extract_dir)

        video_files = list(Path(extract_dir).rglob("*.mp4"))

        # Limit to first 10 videos for the example
        video_files = video_files[:10]

        video_dataset = ray.data.from_items(
            [
                {
                    "video_path": str(video_file),
                    "video_url": f"file://{video_file}",
                    "text": "Describe what happens in this video.",
                }
                for video_file in video_files
            ]
        )

        return video_dataset
    except Exception as e:
        print(f"Error loading dataset: {e}")
        return None

Next, configure the VLM processor with the essential settings:

video_processor_config = vLLMEngineProcessorConfig(
    model_source="Qwen/Qwen3-VL-4B-Instruct",
    engine_kwargs=dict(
        tensor_parallel_size=4,
        pipeline_parallel_size=1,
        trust_remote_code=True,
        limit_mm_per_prompt={"video": 1},
        mm_processor_kwargs={
            "size": {
                "shortest_edge": 65536,
                "longest_edge": 20 * 1088 * 1920,
            },
            "do_sample_frames": False,
        },
    ),
    batch_size=1,
    accelerator_type="L4",
    concurrency=1,
    prepare_multimodal_stage={
        "enabled": True,
        "model_config_kwargs": dict(
            # See available model config kwargs at https://docs.vllm.ai/en/latest/api/vllm/config/#vllm.config.ModelConfig
            allowed_local_media_path="/tmp",
            media_io_kwargs={"video": {"num_frames": 20, "fps": 2}},
        ),
    },
    chat_template_stage=True,
    tokenize_stage=True,
    detokenize_stage=True,
)

Ray Data LLM forwards mm_processor_kwargs to vLLM, which calls the model’s Hugging Face processor with those arguments. The Hugging Face processor defines the accepted keys, and they differ by model family. For example, Qwen2-VL accepts max_pixels and Qwen3-VL accepts size. For the keys your model accepts, see the source of its Hugging Face processor, such as Qwen3VLVideoProcessor.

Note

The following arguments configure multimodal inputs:

engine_kwargs.limit_mm_per_prompt={"video": 1}: Caps the number of videos per request.
engine_kwargs.mm_processor_kwargs.size: Sets the resize budget for each input. The Hugging Face processor resizes each input so that its total pixel count falls between shortest_edge and longest_edge. For a video, the total counts the pixels in all frames.
engine_kwargs.mm_processor_kwargs.do_sample_frames=False: Skips the Hugging Face processor’s own frame sampling because media_io_kwargs already produced the final frames. Set this option whenever frame sampling already happened upstream.
prepare_multimodal_stage.model_config_kwargs.allowed_local_media_path: Required for file:// or local-path media inputs.
prepare_multimodal_stage.model_config_kwargs.media_io_kwargs: Configures frame sampling at decode time.

Warning

If a multimodal input exceeds mm_processor_kwargs.size, the Hugging Face processor’s smart_resize downscales it automatically. Set size.longest_edge to match the largest input you expect to process, which is height * width for an image and num_frames * height * width for a video.

Define preprocessing and postprocessing functions that convert dataset rows into the format the VLM expects and extract the model responses. In the preprocessor, structure video data as part of an OpenAI-compatible message.

def video_preprocess(row: dict) -> dict:
    """
    Preprocessing function for video-language model inputs.

    Converts dataset rows into the format expected by the VLM:
    - System prompt for analysis instructions
    - User message with text and video content
    - Sampling parameters
    - Multimodal processor kwargs for video processing
    """
    return {
        "messages": [
            {
                "role": "system",
                "content": (
                    "You are a helpful assistant that analyzes videos. "
                    "Watch the video carefully and provide detailed descriptions."
                ),
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": row["text"],
                    },
                    {
                        "type": "video_url",
                        "video_url": {"url": row["video_url"]},
                    },
                ],
            },
        ],
        "sampling_params": {
            "temperature": 0.3,
            "max_tokens": 150,
            "detokenize": False,
        },
    }


def video_postprocess(row: dict) -> dict:
    return {
        "resp": row["generated_text"],
    }

Finally, run the VLM inference:

    """Run the complete VLM video example workflow."""
    config = create_vlm_video_config()
    video_dataset = load_video_dataset()

    if video_dataset:
        # Build processor with preprocessing and postprocessing
        processor = build_processor(
            config, preprocess=video_preprocess, postprocess=video_postprocess
        )

        print("VLM video processor configured successfully")
        print(f"Model: {config.model_source}")
        print(f"Has multimodal support: {config.prepare_multimodal_stage.get('enabled', False)}")
        result = processor(video_dataset).take_all()
        return config, processor, result

Run audio batch inference with an omni-modal model#

First, load an audio dataset:

    """
    Load audio dataset from MRSAudio Hugging Face dataset.
    """
    try:
        from datasets import load_dataset
        from huggingface_hub import hf_hub_download
        import base64

        dataset_name = "MRSAudio/MRSAudio"

        dataset = load_dataset(dataset_name, split="train")

        audio_items = []

        # Limit to first 10 samples for the example
        num_samples = min(10, len(dataset))
        for i in range(num_samples):
            item = dataset[i]

            audio_path = hf_hub_download(
                repo_id=dataset_name, filename=item["path"], repo_type="dataset"
            )

            with open(audio_path, "rb") as f:
                audio_bytes = f.read()

            audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
            audio_items.append(
                {
                    "audio_data": audio_base64,
                    "text": item.get("text", "Describe this audio."),
                }
            )

        audio_dataset = ray.data.from_items(audio_items)
        return audio_dataset
    except Exception as e:
        print(f"Error loading dataset: {e}")
        return None

Next, configure the omni-modal processor with the essential settings:

audio_processor_config = vLLMEngineProcessorConfig(
    model_source="Qwen/Qwen2.5-Omni-3B",
    task_type="generate",
    engine_kwargs=dict(
        limit_mm_per_prompt={"audio": 1},
    ),
    batch_size=16,
    accelerator_type="L4",
    concurrency=1,
    prepare_multimodal_stage={
        "enabled": True,
        "chat_template_content_format": "openai",
    },
    chat_template_stage=True,
    tokenize_stage=True,
    detokenize_stage=True,
)

Define preprocessing and postprocessing functions that convert dataset rows into the format the omni-modal model expects and extract the model responses. In the preprocessor, structure audio data as part of an OpenAI-compatible message. You can pass either an audio URL or audio binary data.

"""Supported audio input formats: audio URL, audio binary data"""
{
    "messages": [
        {
            "role": "system",
            "content": "Provide a detailed description of the audio."
        },
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "Describe what happens in this audio."},
                # Option 1: Provide audio URL
                {"type": "audio_url", "audio_url": {"url": "https://example.com/audio.wav"}},
                # Option 2: Provide audio binary data
                {"type": "input_audio", "input_audio": {"data": audio_base64, "format": "wav"}},
            ]
        },
    ]
}

def audio_preprocess(row: dict) -> dict:
    """
    Preprocessing function for audio-language model inputs.

    Converts dataset rows into the format expected by the Omni model:
    - System prompt for analysis instructions
    - User message with text and audio content
    - Sampling parameters
    """
    return {
        "messages": [
            {
                "role": "system",
                "content": "You are a helpful assistant that analyzes audio. "
                "Listen to the audio carefully and provide detailed descriptions.",
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": row["text"],
                    },
                    {
                        "type": "input_audio",
                        "input_audio": {
                            "data": row["audio_data"],
                            "format": "wav",
                        },
                    },
                ],
            },
        ],
        "sampling_params": {
            "temperature": 0.3,
            "max_tokens": 150,
            "detokenize": False,
        },
    }


def audio_postprocess(row: dict) -> dict:
    return {
        "resp": row["generated_text"],
    }

Finally, run the omni-modal inference:

    """Run the complete Omni audio example workflow."""
    config = create_omni_audio_config()
    audio_dataset = load_audio_dataset()

    if audio_dataset:
        # Build processor with preprocessing and postprocessing
        processor = build_processor(
            config, preprocess=audio_preprocess, postprocess=audio_postprocess
        )

        print("Omni audio processor configured successfully")
        print(f"Model: {config.model_source}")
        print(f"Has multimodal support: {config.prepare_multimodal_stage.get('enabled', False)}")
        result = processor(audio_dataset).take_all()
        return config, processor, result

Generate embeddings#

For embedding models, set task_type="embed" and disable chat templating:

    import ray
    from ray.data.llm import vLLMEngineProcessorConfig, build_processor

    embedding_config = vLLMEngineProcessorConfig(
        model_source="sentence-transformers/all-MiniLM-L6-v2",
        task_type="embed",
        engine_kwargs=dict(
            enable_prefix_caching=False,
            enable_chunked_prefill=False,
            max_model_len=256,
            enforce_eager=True,
        ),
        batch_size=32,
        concurrency=1,
        chat_template_stage=False,  # Skip chat templating for embeddings
        detokenize_stage=False,     # Skip detokenization for embeddings
    )

    embedding_processor = build_processor(
        embedding_config,
        preprocess=lambda row: dict(prompt=row["text"]),
        postprocess=lambda row: {
            "text": row["prompt"],
            "embedding": row["embeddings"],
        },
    )

    texts = [
        "Hello world",
        "This is a test sentence",
        "Embedding models convert text to vectors",
    ]
    ds = ray.data.from_items([{"text": text} for text in texts])

    embedded_ds = embedding_processor(ds)
    embedded_ds.show(limit=1)

Embedding models differ from text generation in the following ways:

Use prompt input instead of messages.
Access results through row["embeddings"].

Run classification models#

Ray Data LLM supports batch inference with sequence classification models, such as content classifiers and sentiment analyzers:

import ray
from ray.data.llm import vLLMEngineProcessorConfig, build_processor

# Configure vLLM for a sequence classification model
classification_config = vLLMEngineProcessorConfig(
    model_source="nvidia/nemocurator-fineweb-nemotron-4-edu-classifier",
    task_type="classify",  # Use 'classify' for sequence classification models
    engine_kwargs=dict(
        max_model_len=512,
        enforce_eager=True,
    ),
    batch_size=8,
    concurrency=1,
    chat_template_stage=False,
    detokenize_stage=False,
)

classification_processor = build_processor(
    classification_config,
    preprocess=lambda row: dict(prompt=row["text"]),
    postprocess=lambda row: {
        "text": row["prompt"],
        # Classification models return logits in the 'embeddings' field
        "edu_score": float(row["embeddings"][0])
        if row.get("embeddings") is not None and len(row["embeddings"]) > 0
        else None,
    },
)

# Sample texts with varying educational quality
texts = [
    "lol that was so funny haha",
    "Photosynthesis converts light energy into chemical energy.",
    "Newton's laws describe the relationship between forces and motion.",
]
ds = ray.data.from_items([{"text": text} for text in texts])

if __name__ == "__main__":
    try:
        import torch

        if torch.cuda.is_available():
            classified_ds = classification_processor(ds)
            classified_ds.show(limit=3)
        else:
            print("Skipping classification run (no GPU available)")
    except Exception as e:
        print(f"Skipping classification run due to environment error: {e}")

Classification models differ in the following ways:

Set task_type="classify", or task_type="score" for scoring models.
Set chat_template_stage=False and detokenize_stage=False.
Use direct prompt input instead of messages.
Access classification logits through row["embeddings"].

Query OpenAI-compatible endpoints#

Query deployed models with an OpenAI-compatible API:

    import ray

    OPENAI_KEY = os.environ["OPENAI_API_KEY"]
    ds = ray.data.from_items(["Hand me a haiku."])

    config = HttpRequestProcessorConfig(
        url="https://api.openai.com/v1/chat/completions",
        headers={"Authorization": f"Bearer {OPENAI_KEY}"},
        qps=1,
    )

    processor = build_processor(
        config,
        preprocess=lambda row: dict(
            payload=dict(
                model="gpt-4o-mini",
                messages=[
                    {
                        "role": "system",
                        "content": "You are a bot that responds with haikus.",
                    },
                    {"role": "user", "content": row["item"]},
                ],
                temperature=0.0,
                max_tokens=150,
            ),
        ),
        postprocess=lambda row: dict(
            response=row["http_response"]["choices"][0]["message"]["content"]
        ),
    )

    ds = processor(ds)
    print(ds.take_all())

Configure tokenization disaggregation#

By default, tokenization and detokenization run as separate CPU stages in the processor pipeline. This offloads tokenizer work from the GPU stage, so you can scale the CPU and GPU stages independently.

Note

When you enable the detokenize stage, set detokenize=False in sampling_params so the engine returns raw token IDs for the CPU stage to decode. When you disable it, set detokenize=True so the engine decodes the output itself.

The default disaggregated mode runs tokenize and detokenize as separate CPU stages:

config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={"max_model_len": 4096},
    concurrency=1,
    batch_size=64,
    tokenize_stage=True,
    detokenize_stage=True,
)

processor = build_processor(
    config,
    preprocess=lambda row: dict(
        messages=[
            {"role": "user", "content": row["item"]},
        ],
        sampling_params=dict(
            temperature=0.3,
            max_tokens=250,
            # Let the vLLMEngineProcessor's CPU detokenize stage handle detokenization
            detokenize=False,
        ),
    ),
    postprocess=lambda row: dict(resp=row["generated_text"]),
)

Alternatively, use aggregated mode. Disable these stages so the vLLM engine handles tokenization and detokenization internally:

config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={"max_model_len": 4096},
    concurrency=1,
    batch_size=64,
    tokenize_stage=TokenizerStageConfig(enabled=False),
    detokenize_stage=DetokenizeStageConfig(enabled=False),
)

processor = build_processor(
    config,
    preprocess=lambda row: dict(
        messages=[
            {"role": "user", "content": row["item"]},
        ],
        sampling_params=dict(
            temperature=0.3,
            max_tokens=250,
            # Let the vLLM engine handle detokenization
            detokenize=True,
        ),
    ),
    postprocess=lambda row: dict(resp=row["generated_text"]),
)

Tip

Disaggregated tokenization helps most when the tokenizer is a bottleneck, for example with large vocabularies or long sequences. If the GPU engine is already saturated, the overhead of the extra stages might not pay off.

Use custom tokenizers#

Use this pattern when vLLM supports a model but Hugging Face transformers doesn’t, for example Mistral Tekken (mistral), DeepSeek-V3 (deepseek_v32), or Grok-2 (grok2). The built-in ChatTemplate, Tokenize, and Detokenize stages rely on Hugging Face and fail for these models. The following example disables the built-in CPU stages and replaces them with map_batches callables.

The custom chat template stage converts OpenAI-format messages into the prompt string the model expects. You need this stage because each model family defines its own chat format:

from typing import Any, Dict, List
from vllm.tokenizers import get_tokenizer


class VLLMChatTemplate:
    """Apply a chat template using vLLM's tokenizer."""

    def __init__(self, model_id: str, tokenizer_mode: str = "auto"):
        self.tokenizer = get_tokenizer(
            model_id,
            tokenizer_mode=tokenizer_mode,
            trust_remote_code=True,
        )

    async def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]:
        prompts: List[str] = []
        all_messages: List[List[Dict[str, Any]]] = []

        for messages in batch["messages"]:
            if hasattr(messages, "tolist"):
                messages = messages.tolist()
            all_messages.append(messages)

            add_generation_prompt = messages[-1]["role"] == "user"
            prompt = self.tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=add_generation_prompt,
                continue_final_message=not add_generation_prompt,
            )
            prompts.append(prompt)

        return {
            "prompt": prompts,
            "messages": all_messages,
            "sampling_params": batch["sampling_params"],
        }

The custom tokenize stage converts the prompt string into token IDs for the model:

class VLLMTokenize:
    """Tokenize text prompts using vLLM's tokenizer."""

    def __init__(self, model_id: str, tokenizer_mode: str = "auto"):
        self.tokenizer = get_tokenizer(
            model_id,
            tokenizer_mode=tokenizer_mode,
            trust_remote_code=True,
        )

    async def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]:
        all_tokenized: List[List[int]] = [
            self.tokenizer.encode(prompt) for prompt in batch["prompt"]
        ]

        return {
            "tokenized_prompt": all_tokenized,
            "messages": batch["messages"],
            "sampling_params": batch["sampling_params"],
        }

The custom detokenize stage is optional. It decodes generated token IDs back to text. The vLLM engine already returns generated_text, so you need this stage only for custom decoding, such as different skip_special_tokens settings:

class VLLMDetokenize:
    """Detokenize generated token IDs using vLLM's tokenizer."""

    def __init__(self, model_id: str, tokenizer_mode: str = "auto"):
        self.tokenizer = get_tokenizer(
            model_id,
            tokenizer_mode=tokenizer_mode,
            trust_remote_code=True,
        )

    async def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]:
        decoded: List[str] = []
        for tokens in batch["generated_tokens"]:
            if hasattr(tokens, "tolist"):
                tokens = tokens.tolist()
            decoded.append(self.tokenizer.decode(tokens, skip_special_tokens=True))

        return {
            **batch,
            "generated_text_custom": decoded,
        }

Build a processor with the built-in stages disabled and compose the full pipeline:

MODEL_ID = "unsloth/Llama-3.1-8B-Instruct"

config = vLLMEngineProcessorConfig(
    model_source=MODEL_ID,
    engine_kwargs=dict(
        max_model_len=4096,
        trust_remote_code=True,
        tokenizer_mode="auto",
    ),
    batch_size=4,
    concurrency=1,
    # Disable built-in stages -- we handle them via map_batches.
    chat_template_stage=False,
    tokenize_stage=False,
    detokenize_stage=False,
)

processor = build_processor(
    config,
    postprocess=lambda row: {
        "generated_text": row.get("generated_text", ""),
        "generated_tokens": row.get("generated_tokens", []),
        "num_input_tokens": row.get("num_input_tokens", 0),
        "num_generated_tokens": row.get("num_generated_tokens", 0),
    },
)

ds = ds.map_batches(
    VLLMChatTemplate,
    fn_constructor_kwargs={"model_id": MODEL_ID},
    concurrency=1,
    batch_size=4,
)

ds = ds.map_batches(
    VLLMTokenize,
    fn_constructor_kwargs={"model_id": MODEL_ID},
    concurrency=1,
    batch_size=4,
)

ds = processor(ds)

ds = ds.map_batches(
    VLLMDetokenize,
    fn_constructor_kwargs={"model_id": MODEL_ID},
    concurrency=1,
    batch_size=4,
)

Note

This example uses a standard model because the models that require vLLM’s custom tokenizer are too large for Ray CI environments. The pattern is the same for those models. Replace MODEL_ID and set tokenizer_mode explicitly.

How does Ray Data LLM handle failures?#

The following sections describe how Ray Data LLM handles row failures and actor failures, and how to resume a pipeline from a checkpoint.

Continue processing after row failures#

To enable row-level fault tolerance, set the should_continue_on_error parameter to True in the processor config. If a single row fails because of a request-level error from the engine, the job continues processing the remaining rows. Use this option for long-running jobs where you want to minimize the impact of request failures.

# Row-level fault tolerance configuration
config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    concurrency=1,
    batch_size=64,
    should_continue_on_error=True,
)

What happens when an actor dies?#

When an actor dies during pipeline execution, the actor restarts and rejoins the actor pool to process the remaining rows. This behavior is on by default and needs no extra configuration.

Resume from a checkpoint#

Ray Data supports checkpoint recovery, so you can resume pipeline execution from a checkpoint stored in local or cloud storage. Checkpointing works only for pipelines that start with a read operation and end with a write operation. For checkpointing to take effect, successful blocks must reach the write sink before a failure occurs. After a failure, resume processing from the checkpoint in a later run.

First, set up the checkpoint configuration and specify the ID column for checkpointing.

from ray.data.checkpoint import CheckpointConfig

ctx = ray.data.DataContext.get_current()
ctx.checkpoint_config = CheckpointConfig(
    id_column="id",
    checkpoint_path=checkpoint_path,
    delete_checkpoint_on_success=False,
)

Then, include a read operation and a write operation in the pipeline to enable checkpoint recovery. Preserve the ID column during postprocessing so the checkpoint stores it.

processor_config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={
        "max_num_batched_tokens": 4096,
        "max_model_len": 4096,
    },
    concurrency=1,
    batch_size=16,
)

processor = build_processor(
    processor_config,
    preprocess=lambda row: dict(
        id=row["id"], # Preserve the ID column for checkpointing
        messages=[{"role": "user", "content": row["message"]}],
        sampling_params=dict(
            temperature=0.3,
            max_tokens=10,
        ),
    ),
    postprocess=lambda row: {
        "id": row["id"], # Preserve the ID column for checkpointing
        "answer": row.get("generated_text"),
    },
)

ds = ray.data.read_parquet(input_path)
ds = processor(ds)
ds.write_parquet(output_path)

To resume from a checkpoint, run the same code again. Ray Data discovers the checkpoint and resumes from the last successful block.

Advanced configuration#

The following sections describe options for tuning parallelism, resources, model loading, and throughput.

Configure model parallelism#

For large models that don’t fit on a single GPU, use tensor and pipeline parallelism:

# Model parallelism configuration for larger models
# tensor_parallel_size=2: Split model across 2 GPUs for tensor parallelism
# pipeline_parallel_size=2: Use 2 pipeline stages (total 4 GPUs needed)
# Total GPUs required = tensor_parallel_size * pipeline_parallel_size = 4
config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={
        "max_model_len": 16384,
        "tensor_parallel_size": 2,
        "pipeline_parallel_size": 2,
        "enable_chunked_prefill": True,
        "max_num_batched_tokens": 2048,
    },
    concurrency=1,
    batch_size=32,
    accelerator_type="L4",
)

Configure cross-node parallelism#

Ray Data LLM supports cross-node parallelism, including tensor parallelism and pipeline parallelism. Configure the parallelism level through engine_kwargs. The distributed_executor_backend defaults to "ray" for cross-node support.

config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={
        "enable_chunked_prefill": True,
        "max_num_batched_tokens": 4096,
        "max_model_len": 16384,
        "pipeline_parallel_size": 4,
        "tensor_parallel_size": 4,
        "distributed_executor_backend": "ray",
    },
    batch_size=32,
    concurrency=1,
)

To control how Ray places vLLM engine workers across nodes, customize the placement group configuration. Use bundle_per_worker for a basic per-worker resource specification that’s replicated automatically based on TP*PP, the tensor parallel size times the pipeline parallel size. Use bundles for full control over individual bundles. You can specify the degree of tensor and pipeline parallelism, but the vLLM engine manages the assignment of model ranks to GPUs.

Note

In each bundle dictionary, an omitted CPU or GPU key counts as 0. Specify the resources each worker needs explicitly.

# Simple: specify resources per worker, auto-replicated by TP*PP (4 workers here)
# Alternative: use "bundles": [{"GPU": 1}] * 4 for explicit bundle control
config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={
        "enable_chunked_prefill": True,
        "max_num_batched_tokens": 4096,
        "max_model_len": 16384,
        "pipeline_parallel_size": 2,
        "tensor_parallel_size": 2,
        "distributed_executor_backend": "ray",
    },
    batch_size=32,
    concurrency=1,
    placement_group_config={
        "bundle_per_worker": {"GPU": 1},
        "strategy": "STRICT_PACK",
    },
)

Configure individual stages#

Configure each pipeline stage separately for fine-grained resource control:

config = vLLMEngineProcessorConfig(
    model_source="meta-llama/Llama-3.1-8B-Instruct",
    chat_template_stage={
        "enabled": True,
        "batch_size": 256,
        "concurrency": 4,
    },
    tokenize_stage={
        "enabled": True,
        "batch_size": 512,
        "num_cpus": 0.5,
    },
    detokenize_stage={
        "enabled": True,
        "concurrency": (2, 8),  # Autoscaling pool
    },
)

See stage config classes for all available fields.

Use LoRA adapters#

Configure multi-LoRA batch inference as follows:

# Multi-LoRA configuration
config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={
        "enable_lora": True,
        "max_lora_rank": 32,
        "max_loras": 1,
        "max_model_len": 16384,
    },
    concurrency=1,
    batch_size=32,
)

For details, see the vLLM with LoRA example.

Accelerate model loading with RunAI Model Streamer#

Use RunAI Model Streamer for faster model loading from cloud storage:

Note

To install vLLM with the RunAI dependencies, run pip install -U "vllm[runai]>=0.10.1".

# RunAI streamer configuration for optimized model loading
# Note: Install vLLM with runai dependencies: pip install -U "vllm[runai]>=0.10.1"
config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={
        "load_format": "runai_streamer",
        "max_model_len": 16384,
    },
    concurrency=1,
    batch_size=64,
)

Tune concurrent batch processing#

Two parameters, max_concurrent_batches and max_tasks_in_flight_per_actor, control concurrent batch processing. Understanding how they interact helps you achieve optimal throughput.

What do the parameters control?#

max_concurrent_batches, default: 8: The number of batches that can execute concurrently in a single vLLM engine actor. This overlaps batch processing to hide tail latency. The optimal value depends on the workload.
max_tasks_in_flight_per_actor, default: 2 * max_concurrent_batches: The number of tasks Ray Data can queue per actor before waiting for results. Ray Data uses this queue to prefetch tasks so they’re ready when the actor finishes processing. If you don’t set it, Ray Data uses DataContext.max_tasks_in_flight_per_actor when that’s set, and otherwise twice max_concurrent_batches, which is 16 at the default.

How do the parameters work together?#

These parameters control different parts of the pipeline:

max_tasks_in_flight_per_actor controls how many tasks Ray Data sends to the actor queue.
max_concurrent_batches controls how many batches can execute simultaneously.

When max_tasks_in_flight_per_actor is less than max_concurrent_batches, Ray Data actors aren’t fully saturated. To maximize throughput, increase max_tasks_in_flight_per_actor to keep the actor task queue saturated.

# Tuning concurrent batch processing
# Configure both parameters together for optimal throughput
config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={
        "enable_chunked_prefill": True,
        "max_num_batched_tokens": 4096,
    },
    batch_size=64,
    # Dataset-level concurrency (number of actor replicas)
    concurrency=1,
    # Number of batches that can run concurrently per actor (default: 8)
    max_concurrent_batches=8,
    # Number of tasks Ray Data queues per actor (default: 16)
    # Increase to keep actor task queue saturated
    max_tasks_in_flight_per_actor=16,
)

Troubleshooting#

The following sections describe common issues and how to resolve them.

vLLM compatibility#

The Ray project fully tests each Ray release with a compatible vLLM version. The following table lists the compatible vLLM version for recent Ray releases.

Ray release	vLLM version
nightly	0.30.0
2.59.0	0.27.0
2.58.0	0.26.0
2.57.0	0.25.1
2.56.0	0.22.0
2.55.0	0.18.0
2.54.0	0.15.0
2.53.0	0.12.0
2.52.0	0.11.0
2.51.0	0.11.0
2.50.0	0.10.2

Resolve CUDA out-of-memory errors#

If you hit CUDA out-of-memory errors, try the following strategies:

Reduce the batch size. Start with 8 to 16 and increase it gradually.
Lower max_num_batched_tokens from 4096 to 2048 or 1024.
Decrease max_model_len to use shorter context lengths.
Set gpu_memory_utilization to a value between 0.75 and 0.85 instead of the default 0.90.

# GPU memory management configuration
# If you encounter CUDA out of memory errors, try these optimizations:
config_memory_optimized = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={
        "max_model_len": 8192,
        "max_num_batched_tokens": 2048,
        "enable_chunked_prefill": True,
        "gpu_memory_utilization": 0.85,
        "block_size": 16,
    },
    concurrency=1,
    batch_size=16,
)

# For very large models or limited GPU memory:
config_minimal_memory = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    engine_kwargs={
        "max_model_len": 4096,
        "max_num_batched_tokens": 1024,
        "enable_chunked_prefill": True,
        "gpu_memory_utilization": 0.75,
    },
    concurrency=1,
    batch_size=8,
)

Load models at scale#

On large clusters, Hugging Face might rate-limit model downloads. Cache models to S3 or GCS:

python -m ray.llm.utils.upload_model \
    --model-source facebook/opt-350m \
    --bucket-uri gs://my-bucket/path/to/model

Then reference the remote path in your config:

# S3 hosted model configuration
s3_config = vLLMEngineProcessorConfig(
    model_source="s3://your-bucket/your-model-path/",
    engine_kwargs={
        "load_format": "runai_streamer",
        "max_model_len": 16384,
    },
    concurrency=1,
    batch_size=64,
)

Fix the vLLM NIXL EP dependency incompatibility#

Known issue

If you install Ray and vLLM directly, you might see the following NIXL EP incompatibility error:

ImportError: libcudart.so.12: cannot open shared object file: No such file or directory

Remove the incompatible package, or make sure the installed nixl_ep package is compatible with the CUDA runtime and vLLM build in your environment.

Get help#

If this guide doesn’t cover your issue, use the following resources:

Ray GitHub Issues: Report bugs or request features.
Ray Slack: Get help from the community.
Ray Discourse Forum: Ask questions and share knowledge.
Ray LLM Office Hours: Learn about new Ray LLM features, ask questions, and get guidance from the team.
- Past Office Hours Recordings: View recordings from previous sessions.

Usage data collection#

Ray collects anonymous usage data to improve Ray Data LLM. To opt out, see usage stats collection.

Working with LLMs#

Quickstart: Run batch inference with vLLM#

How does the processor pipeline work?#

Scale to multiple GPUs#

Generate text#

Run batch inference on multimodal data#

Run image batch inference with a vision language model#

Run video batch inference with a vision language model#

Run audio batch inference with an omni-modal model#

Generate embeddings#

Run classification models#

Query OpenAI-compatible endpoints#

Configure tokenization disaggregation#

Use custom tokenizers#

How does Ray Data LLM handle failures?#

Continue processing after row failures#

What happens when an actor dies?#

Resume from a checkpoint#

Advanced configuration#

Configure model parallelism#

Configure cross-node parallelism#

Configure individual stages#

Use LoRA adapters#

Accelerate model loading with RunAI Model Streamer#

Tune concurrent batch processing#

What do the parameters control?#

How do the parameters work together?#

Share a vLLM engine with Ray Serve#

Troubleshooting#

vLLM compatibility#

Resolve CUDA out-of-memory errors#

Load models at scale#

Fix the vLLM NIXL EP dependency incompatibility#

Get help#

Usage data collection#