Enable packing + compile for DPO with wasted tokens metric

[finbarrtimbers]

Summary

• Enable packing + torch.compile for DPO training to reduce padding waste
• Fix cu_seq_lens offset bug when concatenating padded chosen/rejected sequences
• Add tokens_per_second_per_gpu metric to wandb

DPO Packing Performance Results (post MFU-fix)

All runs on finbarr/profile-dpo branch, 2 nodes x 8 H100 GPUs, OLMo2-7B, 16k seq len, activation_memory_budget=0.1.

Batch size comparison (no shard_degree, 16-way FSDP)

Batch Size	Token Count	MFU (last)	TPS	TPS/GPU	sec/step	Experiment
bs=2	57,664	2.46%	9,276	580	6.21s	01KGSYNQ5RYZD1SVWWYDTQNY49
bs=4	58,640	2.51%	9,465	592	6.19s	01KGT1DNDKHZXR1PM2D3ADHD5M
bs=8	431,120	19.33%	69,611	4,351	6.19s	01KGT22ECAJPVPVQP5BRAAA9Q1
bs=16	453,184	19.71%	73,067	4,567	6.19s	01KGTEY8SZ3CWKHBX2KXD7HHAB
bs=32	524,288	22.48%	84,513	5,282	6.19s	01KGTE43M9SBGD9D1Y7Z010EZ3

shard_degree=8 vs no shard_degree (bs=8, full epoch, 57 steps after warmup)

Config	MFU	TPS	TPS/GPU	sec/step	Token Count	Experiment
no shard (16-way FSDP)	8.82% ± 2.54	32,608 ± 9,141	2,038 ± 571	4.212 ± 0.010	233,151 ± 59,605	01KH1S31HR9Y6BRN8MTNY50PV3 / wandb
shard=8, replicas=2 (8-way FSDP)	8.84% ± 2.55	32,670 ± 9,166	2,042 ± 573	4.210 ± 0.009	233,151 ± 59,605	01KH1QG1HYHEYVHRFBXC79MYA3 / wandb

Result: shard_degree=8 and no shard_degree produce identical performance. Over a full epoch (57 steps after warmup), all metrics are within noise. For DPO with 2 nodes, shard_degree has no effect — the default 16-way FSDP is fine.

Binary search: activation_memory_budget with bs=4

Tested both with shard_degree=8 and without shard_degree (16-way FSDP). All budgets > 0.1 OOM on the same 12.25 GiB logits allocation from DPO's concatenated forward pass.

With shard_degree=8 (8-way FSDP):

Budget	Result	Experiment
0.5	OOM (688 MiB alloc failed)	01KH1JHGY6STAF3GCCJTEPTWQS
0.3	OOM (12.25 GiB alloc failed)	01KH1K8WYQ2PA3Y2NZBG2GKPEF
0.2	OOM (12.25 GiB alloc failed)	01KH1KRZ8J56RM81A0RNK6SQJA

Without shard_degree (16-way FSDP):

Budget	Result	Experiment
0.5	OOM (688 MiB alloc failed)	01KH1VFZB95VQF38A35HA7HHPV
0.3	OOM (12.25 GiB alloc failed)	01KH1W3HV9DHGV87X8J6AK1GT8
0.2	OOM (12.25 GiB alloc failed)	01KH1WJ86T4Q3D9RNH171ZPY49
0.15	OOM (12.25 GiB alloc failed)	01KH1X0C2YTCR7K4HT5E9RJ95S

Conclusion: The OOM bottleneck is the logits tensor in DPO's concatenated forward pass (~12.25 GiB = 4 seqs × 2 × 16384 tokens × 152064 vocab × 2 bytes). This allocation is independent of activation checkpointing, so increasing activation_memory_budget beyond 0.1 always fails for bs=4 at 16k seq len. The budget only controls how much recomputation vs caching is used for intermediate activations — it cannot reduce the logits tensor size.

Analysis

• Step time is constant (~6.19s) across all batch sizes (without shard_degree) — the GPU does the same work regardless; the difference is in how many tokens are actually processed per step.
• bs=2/4 waste ~87% of tokens due to poor packing utilization at 16k seq len (only ~58k of 524k theoretical max tokens).
• bs=32 hits 100% packing utilization (524,288 = 16 GPUs × 16,384 tokens × 2 chosen+rejected).
• bs=16 is the sweet spot — 86% packing utilization with minimal sequence truncation (1-4 seqs dropped).
• shard_degree=8 has no effect on throughput — full-epoch comparison shows identical MFU/TPS/sec_per_step.
• MFU and TPS are linearly consistent post-fix (MFU/TPS ratio ~0.000265 across all batch sizes), confirming the sequence counting fix works correctly.
• Note: pre-fix runs (commits before 555d956) reported inflated MFU/TPS because token_count was computed from config (theoretical max) rather than actual tokens processed.

Bug fixes:

1. Fixed cu_seq_lens offset bug (padding_free_collator.py): When concatenating padded chosen/rejected sequences, the offset for rejected cu_seq_lens now correctly uses the padded length instead of actual sequence length.

2. Fixed MFU calculation (olmo_core_callbacks.py): Count actual sequences from batch cu_seq_lens instead of deriving from config, which overcounted when packing truncates sequences.

Pre-filter DPO packing (no more mid-stream truncation)

When DPO packing is enabled, the collator now greedily packs as many complete examples as fit within max_seq_length, rather than packing all examples then truncating mid-stream. This prevents data loss from cutting sequences partway through.

Also fixed the reference cache data loader to use batch_size=1 per rank when packing, ensuring no examples are dropped during cache building.

Runs:

1. Single GPU DPO: Beaker
2. Multi-node DPO: Beaker

[hamishivi]

Generally seems good but some more minor comments. I feel like somehow all this prepro feels too complicated but thats a problem for another day.

[hamishivi]

nit: token counts can probably be longs.

[finbarrtimbers]

Yeah, you're right. The only issue here is that we have to do a manual all-reduce for all of the metrics that we want to calculate the weighted mean for, and that happens in float32 as a bunch of those metrics are float32s:

open-instruct/open_instruct/olmo_core_train_modules.py

Line 173 in 3fb1eff

local_sums_list = [torch.tensor(total_tokens, dtype=torch.float32, device=device)] + [

We could make them all doubles. Idk do you think it's worth refactoring?

[hamishivi]

the logic around return_flash_attn_kwargs is weird to me, I don't like mutable returns. Shouldnt we just always return the cu_seq_lens and let the caller decide if they want to keep or not? Does it add much runtime?

[finbarrtimbers]

no, you're right. I'll do an unconditional return.