feat(training): add DPO training with offline-scored reference model

[Josephasafg]

What does this PR do ?

Adds DPO (Direct Preference Optimization) as a training mode. The reference model's log-probabilities are scored once, offline, and read back during training, so the trainer holds one model and runs at SFT memory and step time.

Changelog

• --mode dpo in scripts/training/run_recipe.py, backed by dpo_train() in training/dpo_train.py: validates the run config and the reference artifact, then hands off to the existing finetune loop.
• Preference-pair dataset, config and loader under data/datasets/preference*.py and data/builders/dpo.py. Rows are TRL-style chat pairs from an HF dataset or a JSONL file. Chosen and rejected rows of a pair always land in the same micro batch, next to each other.
• DPO loss and forward step in training/dpo.py and training/dpo_step.py: sigmoid preference loss, optional SFT term, metrics for preference loss, accuracy, margin and chosen/rejected rewards.
• scripts/dpo/score_reference_logprobs.py, the offline scorer. You run it with the same TP, EP and ETP you'll train with. Extra GPUs beyond that become data-parallel replicas, and rank 0 collects the results and writes the artifact.
• Recipe qwen3_30b_a3b_dpo_config (TP=4, PP=2, EP=4, ETP=1, SP).
• Docs: scripts/dpo/README.md and a DPO section in scripts/training/README.md.

Why the reference is scored offline

The usual DPO trainer keeps a frozen copy of the reference model and runs it every step - a second set of resident weights, an extra forward per step, and the reference is stuck with the policy's parallel layout

Instead, the scorer runs the reference over the dataset once and writes ref_logprobs.jsonl (chosen and rejected completion logprob sums per row) and scoring_metadata.json (dataset, split, tokenizer, sequence length, Parallelism type, prompt key). At startup the trainer compares that metadata against its own config and refuses to run on a mismatch. Without the check a stale or mis-scored artifact does not crash anything, the margins could possibly be wrong.

Pipeline parallelism is not recorded: it only decides which rank runs a layer. Scoring at PP=2 and PP=1 gave the same results over sums on all 8,320 pairs; TP=4 vs TP=8 differed on every pair, so TP is recorded.

The artifact stores plain logprob sums with no beta applied, so one scoring run covers every learning rate and beta you try afterwards.

flowchart LR
    DS[("preference pairs<br/>HF dataset or JSONL")]

    subgraph once["scorer, run once"]
        REF["reference model"]
        ART[/"ref_logprobs.jsonl<br/>chosen / rejected logprob sums per row"/]
        META[/"scoring_metadata.json<br/>dataset, split, tokenizer,<br/>seq len, TP / SP"/]
        REF --> ART
        REF --> META
    end

    subgraph train["trainer, every run"]
        CHK{"metadata matches<br/>run config?"}
        LOADER["pair loader<br/>chosen + rejected adjacent"]
        POL["policy model"]
        LOSS["DPO loss<br/>preference loss, accuracy,<br/>margin, rewards"]
        STEP["backward + optimizer step<br/>finetune loop"]
        CHK -->|"yes"| LOADER
        CHK -->|"no"| STOP["error at startup"]
        LOADER --> POL --> LOSS --> STEP
        LOADER -->|"reference sums"| LOSS
    end

    DS --> REF
    DS -->|"same rows, same order"| LOADER
    ART -->|"joined by row index"| LOADER
    META --> CHK

Loading

Usage

# score the reference once
uv run python -m torch.distributed.run --nproc_per_node=8 scripts/dpo/score_reference_logprobs.py \
    --model Qwen/Qwen3-30B-A3B --dataset HuggingFaceH4/ultrafeedback_binarized --split train_prefs \
    --num-pairs 8192 --max-seq-length 4096 --tp 4 --ep 4 --etp 1 --sequence-parallel \
    --output /data/uf_ref_logprobs

# train
uv run python -m torch.distributed.run --nproc_per_node=8 scripts/training/run_recipe.py \
    --model qwen3_30b_a3b --mode dpo --pretrained_checkpoint /checkpoints/qwen3-30b-a3b \
    dataset.source.path_or_dataset=HuggingFaceH4/ultrafeedback_binarized dataset.source.split=train_prefs \
    dataset.ref_artifact=/data/uf_ref_logprobs dataset.num_pairs=8192

Results

Qwen/Qwen3-30B-A3B with qwen3_30b_a3b_dpo_config at its defaults, 2 nodes x 8 H100 80GB.

Data	HuggingFaceH4/ultrafeedback_binarized train_prefs, 8,192 pairs, one epoch. Validation on 128 pairs of test_prefs every 32 steps
Parallelism	TP=4, PP=2, EP=4, ETP=1, SP, DP=2, distributed optimizer
Batch	64 pairs per step: micro batch 1 pair (2 rows) × 32 accumulation × DP 2; seq length 4096
DPO	beta 0.05, lr 5e-7
Speed	20 s / step mean, 19 s median after warmup (16-34 s range)

Step 1 preference loss is 0.691, i.e. ln 2, as it should be when the policy starts from the reference.

W&B curves from the run above, validation split. The model learns to prefer chosen over rejected

Training preference loss per 32-step window: 0.681, 0.652, 0.636, 0.634.

Accuracy plateaus at 0.68-0.70, in line with published DPO results on UltraFeedback.

Tests

Unit tests, no GPU: tests/unit_tests/data/builders/test_dpo.py, data/datasets/test_preference.py, data/datasets/test_preference_tokenization.py, training/test_dpo.py, training/test_dpo_step.py, training/test_dpo_train.py, scripts/dpo/test_score_reference_logprobs.py, plus cases in training/test_config.py, training/utils/test_pg_utils.py, recipes/qwen/test_qwen3_recipes.py and the scripts/training runner tests.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[cuichenx]

Could you also add at least 1 L1 functional test

[Josephasafg]

@cuichenx Thanks for the great comments. I've updated the PR. PTAL. Thanks!

[Josephasafg]

Could you also add at least 1 L1 functional test

@cuichenx Ive added an L1 test as requested

[cuichenx]

/ok to test 74f8998

[github-actions]

Automatic Claude reviews have been retired. To request a pull-request review, post a comment containing:

/review

Add model=claude to use a Claude reviewer (the default is model=codex). mode=light|strict selects the review depth; for example, /review model=claude mode=strict. Comment /review help for all options.

[cuichenx]

/ok to test 59570d8

[Josephasafg]

@cuichenx Thanks! Seems like two unrelated tests are failing