Implementation:NVIDIA NeMo Curator VLLMEmbedder

Knowledge Sources	NVIDIA NeMo Curator
Domains	NLP, Embeddings, GPU Inference, vLLM
Last Updated	2026-02-14 00:00 GMT

Overview

Provides a high-throughput text embedding stage that leverages vLLM's optimized inference engine for generating embeddings from text data.

Description

VLLMEmbeddingModelStage extends ProcessingStage[DocumentBatch, DocumentBatch] and uses the vLLM LLM class in pooling runner mode to produce text embeddings. Unlike the HuggingFace-based embedding stages, this stage manages its own text input pipeline and does not require a separate tokenizer stage upstream.

The stage initializes a vLLM LLM instance with configurable initialization kwargs, setting sensible defaults for enforce_eager (False), runner ("pooling"), and model_impl ("vllm"). It supports optional pre-tokenization via a HuggingFace AutoTokenizer for finer control over truncation behavior. Model downloads are handled via huggingface_hub.snapshot_download during node setup.

During processing, the stage optionally truncates text by a maximum character count, then calls model.embed() with truncate_prompt_tokens=-1 to generate embeddings. When pre-tokenization is enabled, text is first tokenized with truncation to the model's maximum length, then passed as TokensPrompt objects. Timing metrics are logged for both tokenization and embedding steps.

Usage

Use VLLMEmbeddingModelStage when you want to leverage vLLM's optimized kernel implementations for potentially higher embedding throughput, especially on models that vLLM supports natively. This is an alternative to the HuggingFace-based EmbeddingModelStage and does not need to be paired with a separate TokenizerStage.

Code Reference

Source Location

Repository: NeMo-Curator
File: nemo_curator/stages/text/embedders/vllm.py
Lines: 1-149

Signature

class VLLMEmbeddingModelStage(ProcessingStage[DocumentBatch, DocumentBatch]):
    def __init__(
        self,
        model_identifier: str,
        vllm_init_kwargs: dict[str, Any] | None = None,
        text_field: str = "text",
        pretokenize: bool = False,
        embedding_field: str = "embeddings",
        max_chars: int | None = None,
        cache_dir: str | None = None,
        hf_token: str | None = None,
        verbose: bool = False,
    ): ...

Import

from nemo_curator.stages.text.embedders.vllm import VLLMEmbeddingModelStage

I/O Contract

Inputs

Name	Type	Required	Description
model_identifier	str	Yes	HuggingFace model identifier or local path for the embedding model
vllm_init_kwargs	dict[str, Any] or None	No	Additional keyword arguments passed to the vLLM LLM constructor (default: None)
text_field	str	No	Name of the text column in the input DocumentBatch (default: "text")
pretokenize	bool	No	Whether to pre-tokenize text with a HuggingFace tokenizer before embedding (default: False)
embedding_field	str	No	Name of the output column for embeddings (default: "embeddings")
max_chars	int or None	No	Maximum character count for text truncation (default: None, no truncation)
cache_dir	str or None	No	Directory for caching downloaded model files (default: None)
hf_token	str or None	No	HuggingFace authentication token for gated models (default: None)
verbose	bool	No	Whether to enable verbose logging and progress bars (default: False)

Outputs

Name	Type	Description
DocumentBatch	DocumentBatch	Input data augmented with the original text field and a new embedding column containing embedding vectors

Stage I/O Specification

Method	Returns
inputs()	(["data"], [text_field])
outputs()	(["data"], [text_field, embedding_field])

Usage Examples

Basic Usage

from nemo_curator.stages.text.embedders.vllm import VLLMEmbeddingModelStage

# Create a vLLM embedding stage
embedder = VLLMEmbeddingModelStage(
    model_identifier="intfloat/e5-large-v2",
    text_field="text",
    embedding_field="embeddings",
    max_chars=2000,
)

With Pre-tokenization and Custom vLLM Settings

from nemo_curator.stages.text.embedders.vllm import VLLMEmbeddingModelStage

embedder = VLLMEmbeddingModelStage(
    model_identifier="intfloat/e5-large-v2",
    pretokenize=True,
    vllm_init_kwargs={
        "enforce_eager": True,
        "gpu_memory_utilization": 0.8,
        "max_model_len": 512,
    },
    verbose=True,
)

Implementation Details

vLLM Initialization Defaults

The stage sets the following defaults in vllm_init_kwargs if not explicitly provided:

enforce_eager: False (allows CUDA graph optimizations)
runner: "pooling" (uses the pooling runner for embedding tasks)
model_impl: "vllm" (uses native vLLM model implementation)
download_dir: Set to cache_dir if provided
disable_log_stats: True when not in verbose mode

Resource Requirements

The stage requires 1 CPU and 1 GPU, configured via the Resources dataclass.

Pre-tokenization Flow

When pretokenize is True:

The AutoTokenizer is loaded during setup()
Text is tokenized with batch_encode_plus using truncation to max_model_len
Tokenized inputs are wrapped in TokensPrompt objects from vllm.inputs
These token prompt objects are passed to model.embed() instead of raw text strings

Metrics Logging

The stage tracks and logs timing metrics for both the tokenization step (when pre-tokenization is enabled) and the vLLM embedding step via the _log_metrics method inherited from ProcessingStage.

Related Pages

NVIDIA_NeMo_Curator_EmbedderBase - HuggingFace-based embedding stages and the EmbeddingCreatorStage composite
NVIDIA_NeMo_Curator_ModelStage - Base class for HuggingFace GPU model inference
NVIDIA_NeMo_Curator_TokenizerStage - Standalone tokenizer stage used by the HuggingFace embedding pipeline
Environment:NVIDIA_NeMo_Curator_Python_Linux_Base

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment