DataTrove

To use DataTrove effectively, understand the following hierarchy and concepts:

• Pipeline: A list of processing steps (readers, filters, writers, etc.).
• Executor: The environment that runs a pipeline (e.g., Slurm, local machine, Ray).
• Job: The execution of a pipeline on a specific executor.
• Task: A unit of parallel execution within a job. A job is split into multiple tasks, where each task typically processes a shard of data.
• Shard: A group of input files assigned to a single task.
• File: An individual input file. Note that one file is processed by exactly one task; DataTrove does not automatically split single files into multiple parts. For maximum parallelism, use many medium-sized files rather than one large file.
• Worker: A compute resource (e.g., a CPU core) that executes a task. If you have 50 cores, you can run a LocalPipelineExecutor with workers=50 to process 50 tasks simultaneously.

The vLLM inference benchmark provides performance metrics (tokens per second per GPU) for various models under different configurations. Results are typically obtained through a two-tier optimization sweep on 80GB H100 GPUs.

Key Optimization Levers

• Speculative Decoding (spec): Provides significant speedups for small models (e.g., ngram_6 or suffix_32).
• Batch Size Tuning: Adjusting mns (max number of sequences) and mnbt (max number of tokens) is a consistent way to improve performance across model sizes.
• Tensor Parallelism (tp): Tuning tp is particularly effective for MoE (Mixture of Experts) models.
• GPU Utilization (gmu): Tuning the GPU utilization parameter can impact throughput for certain model scales.

Environment Requirements

To replicate or run these benchmarks, the following environment is used:

• vLLM: 0.15.0
• PyTorch: 2.9.1
• Transformers: 4.57.6
• DataTrove: 0.8.0
• CUDA: 12.8

A pipeline is defined as a list of pipeline blocks. Each block takes a generator of Document objects as input and returns a generator of Document objects.

Common block types include:

• Readers: Read data from various formats and yield Document objects.
• Writers: Save Document objects to disk or cloud.
• Extractors: Extract text from raw formats (e.g., HTML).
• Filters: Remove Document objects based on criteria.
• Stats: Collect statistics on the dataset.
• Tokens: Tokenize data or count tokens.
• Dedup: Perform deduplication tasks.

from datatrove.pipeline.readers import CSVReader
from datatrove.pipeline.filters import SamplerFilter
from datatrove.pipeline.writers import JsonlWriter

pipeline = [
    CSVReader(
        data_folder="/my/input/path"
    ),
    SamplerFilter(rate=0.5),
    JsonlWriter(
        output_folder="/my/output/path"
    )
]

To run a benchmark, create a YAML configuration file. You can define fixed_args that apply to all experiments and an experiments list for specific sweeps. List values in args or fixed_args are automatically expanded into a Cartesian product (parameter sweeps).

Key Configuration Sections:

• script: Path to the inference script.
• continue_on_failure: Boolean to determine if the orchestrator should proceed if a job fails.
• fixed_args: Arguments applied to every experiment.
• experiments: A list of experiment objects. Each can have a name and args. Experiment-level args override fixed_args.

script: "examples/inference/generate_data.py"
continue_on_failure: true

fixed_args:
  qos: "high"
  time: "1:00:00"
  model-max-context: 2048
  max-tokens: 1024
  input-dataset-name: "simplescaling/s1K-1.1"
  input-dataset-split: "train"
  prompt-column: "question"
  output-dataset-name: "s1K-1.1-benchmark"
  output-dir: "data"

experiments:
  - name: "Qwen3-4B"
    args:
      model-name-or-path: "Qwen/Qwen3-4B-Thinking-2507"
      tp: [1, 2, 4]  # Sweep over TP configurations
      speculative-config: [None]

For large-scale production, distribute processing across multiple nodes using Slurm. The script manages a three-part job architecture:

1. inference: A GPU array job that processes data using vLLM and writes Parquet shards to the HF Hub.
2. monitor: A CPU job that updates the dataset card with a live progress bar and ETA.
3. datacard: A CPU job that runs after successful inference to generate a final dataset card with token statistics.

Use --tasks to control the Slurm array size and --workers to specify concurrent jobs.

python examples/inference/generate_data.py \
    --input-dataset-name simplescaling/s1K-1.1 \
    --input-dataset-split train \
    --prompt-column question \
    --model-name-or-path Qwen/Qwen3-4B-Thinking-2507 \
    --output-dataset-name s1K-1.1-dataforge \
    --output-dir data \
    --workers 10 \
    --tasks 20 \
    --examples-per-chunk 50

To use synthetic data generation features, install inference extras using:

uv sync --extra inference

The InferenceRunner supports vLLM, SGLang, OpenAI-compatible HTTPS endpoints, and a local dummy server. It uses asynchronous batching to maintain high GPU utilization.

Custom Rollouts

A rollout is an async callable that receives a Document, a generate(payload) callback, and shared_context kwargs. You can orchestrate multiple sequential or parallel calls within a single rollout.

• Set rollouts_per_document to run the same rollout multiple times per sample; results are stored in document.metadata["rollout_results"].

Recoverable Generation

• Checkpointing: Set checkpoints_local_dir and records_per_chunk to write documents to local chunk files. Failed tasks resume from the last finished chunk. Use ${chunk_index} in the output filename template.
• Deduplication: When checkpointing is enabled, a sqlite-backed RequestCache deduplicates rollouts via payload hashes (requires xxhash and aiosqlite).
• Error Handling: Set skip_bad_requests=True on InferenceRunner to ignore BadRequestError (e.g., context overflows) and continue processing.

from datatrove.data import Document
from datatrove.executor.local import LocalPipelineExecutor
from datatrove.pipeline.inference.run_inference import InferenceConfig, InferenceRunner
from datatrove.pipeline.writers import JsonlWriter

async def simple_rollout(doc: Document, generate):
    payload = {"messages": [{"role": "user", "content": [{"type": "text", "text": doc.text}]}], "max_tokens": 2048}
    return await generate(payload)

documents = [Document(text="What's the weather in Tokyo?", id=str(i)) for i in range(1005)]
config = InferenceConfig(server_type="vllm", model_name_or_path="google/gemma-3-27b-it", rollouts_per_document=1, max_concurrent_generations=500)

LocalPipelineExecutor(
    pipeline=[
        documents,
        InferenceRunner(
            rollout_fn=simple_rollout,
            config=config,
            skip_bad_requests=True,
            records_per_chunk=500,
            checkpoints_local_dir="/fsx/.../translate-checkpoints",
            output_writer=JsonlWriter("s3://.../final_output_data", output_filename="${rank}_chunk_${chunk_index}.jsonl"),
        ),
    ],
    logging_dir="/fsx/.../inference_logs",
    tasks=1,
).run()

Pipelines in DataTrove are platform-agnostic. You can run the same pipeline definition across different environments by using the appropriate PipelineExecutor. To execute a pipeline, call the .run() method on the executor instance.

Common Executor Options

All executors support these options:

• pipeline: A list of the pipeline steps to be executed.
• logging_dir: A directory where logs, statistics, and completion markers are saved. Do not reuse folders for different jobs to avoid overwriting data.
• skip_completed (bool, default True): If True, DataTrove checks the ${logging_dir}/completions folder and skips tasks that have already finished. Set to False to disable this.
• randomize_start_duration (int, default 0): Maximum seconds to delay the start of each task to prevent system overload from simultaneous starts.

from datatrove.executor import LocalPipelineExecutor

executor = LocalPipelineExecutor(
    pipeline=[...],
    logging_dir="logs/",
    tasks=10,
    workers=5
)
executor.run()