What Is DoReMi and How to Use It in Nanotron

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#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#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#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#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#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).

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 with uniform sampling to generate the baseline checkpoint.

4. Run Proxy Training: Launch 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:

# 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:

# 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:

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:

# 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 uses these indices to respect the current weight vector when constructing training batches. Refer to the folder layout requirements in 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}