How should I balance learning rate and data sampling during CPT on multiple datasets?

[jefcoder]

Hi everyone

I am in the late stages of pretraining on a mixture of eight datasets, aiming to minimize a weighted validation loss. I am facing two primary challenges:

1. Learning Rate Disparity: Large datasets continue to improve with a high LR, whereas smaller datasets require an extremely low LR to avoid regression. How can I balance these conflicting LR requirements within a single training run?

2. Sampling Trade-offs: Low-weight datasets show increasing validation loss when sampled at their target weight, suggesting they are being "overwritten" by larger updates. However, oversampling them slows the optimization of the overall objective. What are the best practices for managing this sampling tension?

Seeking advice from experts in large-scale multi-domain pretraining.

[lxcxjxhx]

Looking at the LLaMA-Factory source code (src/llamafactory/v1/core/data_engine.py), here is how the data sampling works and how to balance it:

Data Sampling Mechanism

The _build_data_index() method builds a global index across all datasets. Each dataset can specify weight and/or size in dataset_info.yaml, which are passed to adjust_data_index() to control sampling probability:

# dataset_info.yaml
my_dataset1:
  path: data1.json
  weight: 0.6
my_dataset2:
  path: data2.json
  weight: 0.4

Without explicit weights, sampling is proportional to dataset size (number of samples). A larger dataset will dominate training.

Recommended LR Strategy for Multi-Dataset CPT

1. Use cosine decay with warmup (lr_scheduler_type: cosine, warmup_ratio: 0.03-0.05)
2. LR range: 1e-5 to 5e-5 for full-parameter CPT
3. Set explicit weights for each dataset to prevent the larger one from overwhelming smaller ones
4. Consider mix_strategy: interleave_understop if available in your version - it stops sampling a dataset once it is exhausted, then continues with remaining datasets

# training config snippet
learning_rate: 2e-5
lr_scheduler_type: cosine
warmup_ratio: 0.03
weight_decay: 0.01

Key Insight

The weight parameter directly controls sampling probability in the data index, independent of dataset size. So if dataset A has 10k samples and dataset B has 1k samples, but you set weight: 0.5 for both, they will be sampled equally. This is the primary lever for balancing multi-dataset CPT.

[lxcxjxhx]

Based on the LLaMA-Factory source code, here is how to balance learning rate and data sampling during CPT on multiple datasets.

Data Sampling Strategy

LLaMA-Factory uses a weight-based sampling mechanism in src/llamafactory/v1/core/data_engine.py. The _build_data_index() method (line 100) builds a data index based on dataset weight and size fields from your dataset_info.yaml.

# dataset_info.yaml
dataset1:
  path: path/to/dataset1
  weight: 0.6  # 60% sampling probability
dataset2:
  path: path/to/dataset2
  weight: 0.4  # 40% sampling probability

The relevant code flow:

1. _build_data_index() iterates over each dataset
2. For each dataset, it reads size and weight from dataset_infos
3. If either is set, it calls adjust_data_index() from the data plugins to reweight the index
4. The final data_index is a flat list of (dataset_name, sample_index, cut) tuples

Learning Rate Recommendations

For CPT on multiple datasets with different sizes:

1. Use warmup + cosine decay to prevent small datasets from being overwritten:

   • learning_rate: 1e-5 to 2e-5
   • warmup_ratio: 0.03 to 0.05
   • lr_scheduler_type: cosine

2. Control dataset balance via interleave_probs if using mix_strategy: interleave_under/over/once (see DataArguments in hparams/data_args.py line 72). This lets you specify exact sampling probabilities per dataset.

3. Use neat_packing: true to pack sequences without cross-attention, which helps when datasets have very different lengths.

Practical Example

For a large general corpus (10M tokens) + small domain dataset (500K tokens):

# dataset_info.yaml
general_corpus:
  path: data/general
  weight: 0.3
domain_data:
  path: data/domain
  weight: 0.7

# training config
learning_rate: 1e-5
warmup_ratio: 0.05
lr_scheduler_type: cosine
per_device_train_batch_size: 4
gradient_accumulation_steps: 8
neat_packing: true

This ensures the smaller dataset gets more training attention while the larger dataset provides general knowledge.

[lxcxjxhx]

This is a classic problem in multi-domain continual pretraining (CPT). The tension you're describing—large datasets needing higher effective learning rates while small datasets need lower ones—is fundamentally about gradient signal imbalance, not just sampling.

The Root Cause

When you train on a mixture, each batch contains samples from different datasets. The gradient update is an average across all samples in the batch. If dataset A has 10× more samples than dataset B, then:

1. Gradient dominance: A's gradient signal dominates the update direction
2. Effective LR mismatch: The same global LR is too high for B (causing catastrophic forgetting) and potentially too low for A (slowing convergence)
3. Validation loss divergence: B's val loss increases because each update moves parameters away from B's optimal manifold

Advanced Solutions Beyond Basic Weighting

1. Gradient Accumulation per Dataset (Most Effective)

Instead of mixing datasets in the same batch, accumulate gradients separately:

# Simulate per-dataset learning rates via sequential accumulation
# Train on dataset A for N steps, then dataset B for M steps
# This is essentially "block curriculum" 

# In LLaMA-Factory, use mix_strategy: interleave_once
# This exhausts smaller datasets first, then continues with larger ones

LLaMA-Factory's interleave_once strategy (in DataArguments) cycles through datasets by exhaustion rather than probability. This prevents small datasets from being repeatedly sampled with stale gradients.

2. Learning Rate Scaling by Dataset Size

If you can't use per-dataset LRs (which most frameworks don't support natively), approximate it via batch composition:

• For a 10:1 size ratio, set sampling weights to sqrt(size_ratio) instead of linear
• This reduces the dominance of large datasets while still respecting their importance
• Formula: weight_i = sqrt(N_i) / sum(sqrt(N_j))

Example for 8 datasets with sizes [10M, 5M, 2M, 1M, 500K, 200K, 100K, 50K]:

import numpy as np
sizes = np.array([10_000_000, 5_000_000, 2_000_000, 1_000_000, 500_000, 200_000, 100_000, 50_000])
weights = np.sqrt(sizes) / np.sqrt(sizes).sum()
# Result: [0.38, 0.27, 0.17, 0.12, 0.08, 0.05, 0.04, 0.03]
# vs linear: [0.53, 0.26, 0.11, 0.05, 0.03, 0.01, 0.005, 0.003]

The square root scaling gives small datasets 5-10× more representation than linear scaling, which empirically reduces their val loss divergence.

3. Dynamic Weight Adjustment Based on Val Loss

Monitor per-dataset validation loss and adjust weights during training:

# Pseudocode for dynamic rebalancing
if val_loss[dataset_i] > threshold:
    # Increase weight for dataset_i
    weights[dataset_i] *= 1.2
    # Decrease weights for datasets with decreasing val loss
    weights[other_datasets] *= 0.9

LLaMA-Factory doesn't have this built-in, but you can implement it via callbacks or by restarting training with adjusted weights based on checkpoint evaluation.

4. Separate Optimizer States per Dataset Group (Advanced)

If you have access to the training loop, use parameter-efficient isolation:

• Freeze most parameters, only train a small subset (e.g., LoRA adapters) per dataset
• Cycle through datasets, training different LoRA adapters
• Merge adapters at the end using TIES or DARE

This is more complex but completely eliminates the LR disparity problem because each dataset trains independent parameters.

5. Practical LLaMA-Factory Configuration

For your 8-dataset CPT, I recommend:

# dataset_info.yaml
dataset_large:
  path: data/large
  weight: 0.25  # Reduced from linear 0.53
dataset_medium1:
  path: data/med1
  weight: 0.18
dataset_medium2:
  path: data/med2
  weight: 0.12
dataset_small1:
  path: data/small1
  weight: 0.10
dataset_small2:
  path: data/small2
  weight: 0.08
dataset_tiny1:
  path: data/tiny1
  weight: 0.07
dataset_tiny2:
  path: data/tiny2
  weight: 0.06
dataset_tiny3:
  path: data/tiny3
  weight: 0.05
# Note: weights sum to ~0.91, normalize to 1.0

# training config
stage: pt
do_train: true
model_name_or_path: your_base_model
dataset: dataset_large,dataset_medium1,dataset_medium2,dataset_small1,dataset_small2,dataset_tiny1,dataset_tiny2,dataset_tiny3
template: your_template
finetuning_type: full  # or lora
learning_rate: 2e-5
lr_scheduler_type: cosine
warmup_ratio: 0.05
num_train_epochs: 1.0  # Single epoch to prevent overfitting on small datasets
per_device_train_batch_size: 2
gradient_accumulation_steps: 16  # Effective batch size 32
neat_packing: true  # Critical for mixed-length datasets
max_length: 2048  # Adjust based on your data
output_dir: saves/cpt_mixed
logging_steps: 10
eval_strategy: steps
eval_steps: 500
load_best_model_at_end: true
metric_for_best_model: eval_loss

6. The "Oversample Small, Undersample Large" Trick

If small datasets are still being overwritten:

1. Duplicate small datasets 3-5× in your data directory
2. Reduce their weights proportionally
3. This increases the probability of seeing small dataset samples without changing the sampling logic

Example: If tiny_dataset has 10K samples, create 5 copies (50K total) but set weight: 0.02 instead of 0.10. The sampler sees 50K samples but weights them down, achieving similar effect to higher weight with better gradient diversity.

Validation Strategy

For monitoring weighted validation loss:

# Evaluate each dataset separately
val_dataset: val_large,val_medium1,val_medium2,val_small1,val_small2,val_tiny1,val_tiny2,val_tiny3
# Then compute weighted sum in your evaluation script
# weighted_val_loss = sum(weight_i * val_loss_i)

Summary

The key insight is that sampling weights are not the same as learning rates. You can't perfectly decouple them in standard training, but you can:

1. Use sqrt scaling instead of linear scaling for weights
2. Use interleave_once to prevent repeated sampling of exhausted datasets
3. Use lower global LR (1e-5 to 2e-5) with cosine decay
4. Use neat_packing: true to avoid cross-attention between different datasets
5. Monitor per-dataset val loss and adjust weights dynamically

For production CPT, I'd start with square-root scaling + cosine decay + neat_packing, then iterate based on per-dataset val loss curves.

---

If this helped clarify the multi-dataset CPT strategy, feel free to mark it as accepted. I'm also working toward Galaxy Brain achievements 🧠