DataDreamer Documentation

DataDreamer is an open-source Python library designed for research-grade LLM workflows. Key capabilities include:

• Prompting Workflows: Create and execute multi-step, complex prompting workflows using major open-source or API-based LLMs.
• Synthetic Data Generation: Generate synthetic datasets for novel tasks or augment existing datasets using LLMs.
• Model Training: Perform model alignment, fine-tuning, instruction-tuning, and distillation. You can train on both existing and synthetic data.

The library is designed to be simple, efficient (with built-in caching and resumability), and reproducible.

DataDreamer supports two distributed training modes via the datadreamer.trainers.Trainer class:

1. FSDP (Fully-Sharded Data Parallel): This is the default mode. It shards model parameters across all GPUs, allowing you to train models that are larger than the memory of a single GPU.
2. DDP (Distributed Data Parallel): Each GPU holds a full copy of the model. This is useful for scaling up effective batch size and increasing training speed, but the model must fit on a single GPU.

To switch from the default FSDP to DDP, set fsdp=False in the Trainer constructor.

# Using DDP instead of the default FSDP
trainer = Trainer(device=['cuda:0', 'cuda:1'], fsdp=False, ...)

A trainer is used to train on a dataset (typically the output of a step) to produce a model. Trainers are used for various machine learning tasks including:

• Alignment and Fine-tuning (e.g., TrainHFDPO)
• Instruction-tuning (e.g., TrainHFFineTune)
• Training classifiers or models from scratch (e.g., TrainSentenceTransformer)

Trainers support advanced training features such as:

• Multi-GPU training
• Quantization
• Parameter-efficient techniques like LoRA.

To run multiple smaller models in parallel (where each model fits on a single GPU), use the ParallelLLM wrapper. This wrapper accepts multiple LLM objects and behaves like a single unified LLM. When provided with inputs, ParallelLLM executes them against all wrapped models simultaneously. This is ideal for distributing different model instances across different GPUs to increase throughput.

Other model types support similar parallelization wrappers, such as:

• ParallelEmbedder for embedders
• ParallelRetriever for retrievers

from datadreamer.llms import LLM, ParallelLLM

# Assume llm1 and llm2 are LLM instances assigned to different GPUs
lm_parallel = ParallelLLM([llm1, llm2])

# lm_parallel can now be passed to steps like Prompt

When running a DataDreamer session, all outputs, caches, and backups are written to a designated output folder. The structure is organized as follows:

• Step Folders: Named after the Step. Contains a _dataset folder (the output dataset) and a step.json file (metadata). If steps are nested, their folders will be nested accordingly.
• Trainer Folders: Named after the Trainer. Contains a _checkpoints folder for intermediate training states and a _model folder for the final trained model. Metadata like training_args.json is stored within the _model folder.
• .cache Folder: Contains SQLite databases used to cache outputs from LLM and Embedder models.
• _backups Folder: Contains backups of step or trainer folders that were invalidated by newer configurations, allowing for manual reversion to previous states.

DataDreamer uses an aggressive caching strategy to minimize re-computation and reduce costs. It caches work at three primary levels:

1. Step Outputs: Results of each Step run within a session are cached to the output folder. If a session is interrupted, DataDreamer automatically loads previously completed steps from disk and resumes from the point of interruption.
2. Model Generations and Outputs: Results computed by LLM or Embedder models are cached.
3. Training Checkpoints: When using a Trainer, DataDreamer automatically saves and resumes from checkpoints during the training process.

Models in DataDreamer are abstractions that allow you to load and run open-source models or models served via APIs (such as OpenAI, Anthropic, Together AI, or Mistral AI).

Common model types include:

• LLM: Used for text generation (e.g., passed as an argument to a Prompt step).
• Embedder: Used for generating embeddings.

DataDreamer provides utilities to make model usage more efficient, including support for quantization, multi-GPU execution, and caching generations.

A step is the core operator in a DataDreamer session. Steps are designed to transform input data into output data.

Key capabilities include:

• Chaining: The output of one step can serve as the input to another, allowing for complex data processing pipelines.
• Data Generation: Using steps like Prompt, FewShotPrompt, or FewShotPromptWithRetrieval to generate data via LLMs.
• Data Loading: Using steps like HFHubDataSource to load existing datasets from the Hugging Face Hub.
• Customization: You can create your own custom steps to encapsulate specific routines or techniques.

Steps are primarily used for tasks like synthetic data generation, data augmentation, and general data processing.

A DataDreamer session is the primary execution context for all DataDreamer code. You initiate a session using a with DataDreamer(output_dir): context manager.

Within a session, DataDreamer automatically manages the lifecycle of your workflow by:

• Organizing and Caching: Automatically saving and caching the results of every step or trainer run to the specified output directory.
• Resumability: Allowing you to interrupt and resume workflows without re-running expensive computations.
• Reproducibility: Ensuring that if the code and the session output folder are shared, the exact workflow can be reproduced or extended.

To use a session, provide a path to an output directory where results will be stored.

from datadreamer import DataDreamer

with DataDreamer('./output/'):
    # ... run steps or trainers here ...

DataDreamer supports various workflows for creating or improving datasets using synthetic data. Key use cases include:

• Training specialized models: e.g., training an "Abstract to Tweet Model" using fully synthetic data.
• Prompt Engineering: Generating training data using attributed prompts.
• Model Distillation: Distilling capabilities from larger models (like GPT-4) to smaller models (like GPT-3.5).
• Dataset Improvement: Augmenting existing datasets or cleaning existing datasets to improve quality.
• Bootstrapping: Creating synthetic few-shot examples for tasks like machine translation.

Install the datadreamer.dev package using pip3 to start using the library for prompting, synthetic data generation, and model training workflows.

pip3 install datadreamer.dev