# What Is DoReMi and How to Use It in Nanotron

> Learn about DoReMi, the Nanotron algorithm that auto-optimizes domain mixing. Discover how DoReMi dynamically adjusts training weights for efficient data domain focus and improved model performance.

- Repository: [Hugging Face/nanotron](https://github.com/huggingface/nanotron)
- Tags: how-to-guide
- Published: 2026-03-03

---

**DoReMi is a training algorithm in the huggingface/nanotron framework that dynamically optimizes domain mixing weights by monitoring per-domain losses from a reference model, automatically focusing training on the hardest data domains.**

DoReMi (Domain Reweighting with Minimized Loss) eliminates the need for manual data mixture tuning when training large language models on heterogeneous corpora. Unlike static sampling strategies, this implementation—located under `examples/doremi/`—adjusts sampling probabilities on-the-fly based on real-time difficulty estimates, ensuring optimal allocation of model capacity across domains.

## How DoReMi Optimizes Domain Weights

The algorithm treats per-domain loss as a proxy for **domain difficulty**. During training, DoReMi maintains a reference model trained with uniform sampling to establish baseline loss expectations. The **DoReMiTrainer** then compares current training losses against this reference, increasing sampling weights for domains with higher relative loss (indicating greater difficulty) while applying smoothing to prevent erratic updates.

This reweighting mechanism relies on distributed communication across data-parallel groups via `dist.all_reduce` with the `AVG` operation to aggregate loss statistics. The **smoothing_param** and **step_size** hyperparameters control the convergence speed of the weight vector, preventing over-correction while allowing the distribution to adapt to the model's evolving capabilities.

## Core Components of the DoReMi Implementation

The Nanotron implementation consists of four coordinated components that handle configuration, state management, and distributed training orchestration.

### DoReMiContext

The `DoReMiContext` class, defined in [[`examples/doremi/doremi/doremi_context.py`](https://github.com/huggingface/nanotron/blob/main/examples/doremi/doremi/doremi_context.py)](https://github.com/huggingface/nanotron/blob/main/examples/doremi/doremi/doremi_context.py#L12-L38), serves as the central state container. It stores the list of domain names, the current weight vector (defaulting to uniform `1/num_domains`), and the complete history of weight updates. Key methods include `add_weight_with_history` for logging updates and `get_domain_name` for domain resolution.

### DoReMiConfig

Configuration validation and parsing are handled by `DoReMiConfig` in [[`examples/doremi/doremi/config.py`](https://github.com/huggingface/nanotron/blob/main/examples/doremi/doremi/config.py)](https://github.com/huggingface/nanotron/blob/main/examples/doremi/doremi/config.py#L9-L38). This dataclass ingests user-provided domain lists, optional initial weights, smoothing parameters, step sizes, and the critical **ref_model_resume_checkpoint_path** pointing to the reference model checkpoint. It validates that domain names and reference checkpoints are present before training begins.

### DoReMiTrainer

The proxy training loop is implemented in `DoReMiTrainer`, a subclass of `DistributedTrainer` located at [[`examples/doremi/doremi/trainer.py`](https://github.com/huggingface/nanotron/blob/main/examples/doremi/doremi/trainer.py)](https://github.com/huggingface/nanotron/blob/main/examples/doremi/doremi/trainer.py#L36-L77). This trainer loads the reference model, synchronizes weight tensors across ranks, logs per-domain losses, and saves weight history after each step. It optionally streams metrics to Weights & Biases for monitoring. The core update logic invokes `self.doremi_context.add_weight_with_history` after aggregating losses from all data-parallel workers.

### ReferenceTrainer

To generate the required baseline, `ReferenceTrainer` ([[`examples/doremi/doremi/trainer.py`](https://github.com/huggingface/nanotron/blob/main/examples/doremi/doremi/trainer.py)](https://github.com/huggingface/nanotron/blob/main/examples/doremi/doremi/trainer.py#L200-L207)) executes the initial training phase using fixed uniform or user-specified sampling distributions. This creates the checkpoint that `DoReMiTrainer` loads for loss comparison.

### DistributedSamplerForDoReMi

The data pipeline integrates a specialized sampler in [[`examples/doremi/doremi/dataloader.py`](https://github.com/huggingface/nanotron/blob/main/examples/doremi/doremi/dataloader.py)](https://github.com/huggingface/nanotron/blob/main/examples/doremi/doremi/dataloader.py#L131-L166) that respects the current `domain_weights` when constructing batches. During the proxy run, sampling remains uniform for logging purposes, while the weights are recorded for subsequent analysis and application.

## Step-by-Step Workflow for DoReMi Training

Implementing DoReMi requires a two-stage training process followed by weight extraction.

1. **Prepare the Dataset**: Tokenize your corpus ensuring each sample contains a `domain_ids` column mapping to a domain index. Organize files according to the folder layout specified in [[`examples/doremi/README.md`](https://github.com/huggingface/nanotron/blob/main/examples/doremi/README.md)](https://github.com/huggingface/nanotron/blob/main/examples/doremi/README.md).

2. **Configure DoReMi**: Create a YAML configuration listing domain names, smoothing parameters, step size, and the reference checkpoint path.

3. **Train the Reference Model**: Execute [`train_reference.py`](https://github.com/huggingface/nanotron/blob/main/train_reference.py) with uniform sampling to generate the baseline checkpoint.

4. **Run Proxy Training**: Launch [`train_doremi.py`](https://github.com/huggingface/nanotron/blob/main/train_doremi.py) to perform the dynamic reweighting phase. The trainer outputs `doremi_domain_weights_<step>.pt` files containing the weight history.

5. **Derive Final Weights**: Load the saved history files, average the stored `domain_weights` tensors across steps, and use these optimized values for training your full-scale model.

## Configuration and Code Examples

### Minimal DoReMi Configuration

The proxy training requires a configuration file specifying domains and the reference checkpoint:

```yaml

# examples/doremi/configs/config_280m_llama_proxy.yaml

doremi:
  domain_names: "domain_a,domain_b,domain_c"
  smoothing_param: 0.001
  step_size: 1.0
  ref_model_resume_checkpoint_path: "checkpoints/reference-280m"

```

### Executing the Two-Stage Training

Train the reference model first, then run the DoReMi proxy:

```bash

# Stage 1: Train reference model with uniform sampling

CUDA_DEVICE_MAX_CONNECTIONS=1 torchrun --nproc_per_node=4 \
  examples/doremi/train_reference.py --config-file examples/doremi/configs/config_280m_llama.yaml

# Stage 2: Train proxy model with dynamic reweighting

CUDA_DEVICE_MAX_CONNECTIONS=1 torchrun --nproc_per_node=4 \
  examples/doremi/train_doremi.py --config-file examples/doremi/configs/config_280m_llama_proxy.yaml

```

### Processing Saved Domain Weights

Extract and average the optimized weights from the checkpoint history:

```python
import torch

# Load history checkpoint written by DoReMiTrainer

history_path = "checkpoints/doremi/proxy-280m-llama/doremi_domain_weights_100000.pt"
history = torch.load(history_path)  # List of dicts with "step" and "weight" keys

# Compute average weight vector across training steps

total = sum(entry["weight"] for entry in history)
avg_weights = total / len(history)  # Tensor of shape (num_domains,)

print("Optimized domain weights:", avg_weights)

```

### Applying Optimized Weights to Large-Scale Training

Use the averaged weights in your production configuration:

```yaml

# examples/doremi/configs/config_2.8b_llama_with_tuned_weights.yaml

doremi:
  domain_names: "domain_a,domain_b,domain_c"
  domain_weights: "0.42,0.35,0.23"  # Derived from avg_weights.tolist()

  smoothing_param: 0.001
  step_size: 1.0
  ref_model_resume_checkpoint_path: "checkpoints/reference-2.8b"

```

## Summary

- **DoReMi** dynamically adjusts domain sampling probabilities based on per-domain loss differentials against a reference model.
- The implementation requires two distinct training phases: **reference training** (uniform sampling) and **proxy training** (dynamic reweighting via `DoReMiTrainer`).
- Core classes include `DoReMiContext` for state management, `DoReMiConfig` for validation, and specialized trainers orchestrating the pipeline in `examples/doremi/`.
- Optimized weights are derived by averaging the saved weight history from proxy training checkpoints.
- The technique integrates with Nanotron's distributed training infrastructure, utilizing `dist.all_reduce` for cross-rank loss aggregation.

## Frequently Asked Questions

### What is the purpose of the reference model in DoReMi?

The **reference model** provides a baseline loss expectation for each domain under uniform sampling. By comparing the proxy model's current per-domain loss against this reference, DoReMi identifies which domains are relatively harder or easier for the model to learn, enabling data-driven reweighting rather than relying on manual heuristics.

### How does DoReMi determine which domains are "harder"?

Domains exhibiting **higher loss relative to the reference model** are considered harder and receive increased sampling weights. The algorithm aggregates these loss differentials across all data-parallel ranks using `dist.all_reduce`, then applies a smoothed update rule controlled by the `step_size` and `smoothing_param` configuration parameters to gradually shift the sampling distribution.

### Can I use DoReMi with existing tokenized datasets?

Yes, provided your dataset includes a **`domain_ids`** column that maps each sample to a specific domain index. The `DistributedSamplerForDoReMi` in [`examples/doremi/doremi/dataloader.py`](https://github.com/huggingface/nanotron/blob/main/examples/doremi/doremi/dataloader.py) uses these indices to respect the current weight vector when constructing training batches. Refer to the folder layout requirements in [`examples/doremi/README.md`](https://github.com/huggingface/nanotron/blob/main/examples/doremi/README.md) for specific formatting guidelines.

### Where are the optimized domain weights stored during training?

The `DoReMiTrainer` saves weight history to checkpoint files named **`doremi_domain_weights_<step>.pt`** at regular intervals. These files contain serialized dictionaries with step numbers and weight tensors, which you can load using `torch.load()` to compute average weights for subsequent training runs.