Online Data Mixing

This library contains plugin for online dynamic reward (learnable) based data mixing framework that operates on dynamically mixing datasets online during training while being adapted based on the signals (e.g. training loss, gradnorm etc) from training.

Plugins

Design

Usage in Custom Training Loop

OnlineMixingDataset can be imported easily and integrated into existing training loops with minimal changes. A sample custom training loop implementation can be found here. Given code sample uses two instruction tuning datasets and trains ibm-granite/granite-3.1-2b-instruct model for next token prediction task.

Automatic Categorization

When only a single dataset (without category splits) is passed, the dataset will be embedded with a sentence-transformer model and clustered (K-Means by default) to build pseudo categories used by the online data mixer.

from datasets import load_dataset
from fms_acceleration_odm import OnlineMixingDataset

dataset = load_dataset("tatsu-lab/alpaca", split="train[:1%]")
collator = ...  # e.g., DataCollatorForLanguageModeling(...)

odm_dataset = OnlineMixingDataset(
    dataset_dict=dataset,
    collators_dict={"train": collator},
    eval_dataset_dict={},
    eval_collators_dict={},
    auto_categorize_config={
        "input_column": "text",
        "num_categories": 6,
        "model_name": "sentence-transformers/all-MiniLM-L6-v2",
    },
)

Without an explicit num_categories, a heuristic based on the square root of the dataset size is used. Additional knobs such as category_prefix, batch_size, or clustering-specific kwargs can also be provided through auto_categorize_config.

Metrics

All metrics related to the online data mixing will be logged to odm.jsonl file in the checkpoint output directory.

Rewards

Below are the currently available rewards and we are constantly looking to improve the existing rewards and also add new ones. Further, we encourage users to identify rewards that can help their usecases.

LEARNABILITY and COMBINED require passing templated_eval_dataset_dict (and templated_eval_collators_dict) to OnlineMixingDataset — a few-shot-templated version of eval_dataset_dict, keyed by the same category names, used as the "few-shot" side of the learnability comparison. Constructing OnlineMixingDataset with one of these reward types but without templated_eval_dataset_dict raises a ValueError.

Adding a Custom Reward

Custom rewards can be added to the compute_reward function and adding it to the Reward enum. If the custom reward requires specific set of information from the training loop then _extract_information_from_state_for_reward function has to be extended for extracting such information from trainer state. This is member function of OnlineMixingDataset.

Planned TODOs

Please see issue #153.