amazon-science/dlm-scaling

amazon-science/dlm-scaling

Description: Code release for the paper Revisiting the 16x Pretraining Compute Gap between Diffusion and Autoregressive Language Models

Language: Python

License: NOASSERTION

Stars: 2

Forks: 0

Open issues: 0

Created: 2026-10-09T02:32:16Z

Pushed: 2026-10-14T22:59:19Z

Default branch: main

Fork: no

Archived: no

README:

Revisiting the 16× Pretraining Compute Gap between Diffusion and Autoregressive Language Models

TL;DR: the widely cited 16× compute gap between diffusion and autoregressive language models largely reflects an outdated scaling-law protocol. Under current practice, the gap falls to between 3× and 8×.

Left: the compute multiple *c* = *C*DLM/*C*ARM measures how many times more training compute a diffusion language model (DLM) needs than an autoregressive model (ARM) to reach the same compute-optimal negative log-likelihood (NLL). Right: *c* as the evaluation (ELBO → DUEL), the model design (MDM → BD3LM → E2D2) and the block size are revised.

This repository is the code release behind the paper *Revisiting the 16× Pretraining Compute Gap between Diffusion and Autoregressive Language Models*. Our contribution can be summarized as follows:

• We revisit the scaling laws of DLMs under a protocol that reflects current practice: exact

likelihood evaluation, block diffusion, an encoder–decoder architecture, and a sweep over the block sizes used in practice.

• Pretraining DLMs is far cheaper than previously assumed. Exact likelihood evaluation reduces

the compute multiple to 12.8×; block diffusion with an encoder–decoder architecture reduces it to 7.8× at block size 32; and smaller blocks narrow it to 4.6× at block size 8 and 3.0× at block size 4.

• We find a simple empirical rule for the optimal decoder depth of encoder–decoder block

diffusion, which depends on the block size alone; the gap between variants persists at a held-out budget three times larger than our largest fitted one; and likelihood improvements carry over to downstream benchmark accuracy.

> Reproducibility note. This code is being released solely for academic and scientific > reproducibility purposes, in support of the methods and findings described in the associated > publication. Pull requests are not being accepted in order to maintain the code exactly as it > was used in the paper.

What Is Here

We include the scaling law analysis and benchmark evaluation code for ARM, MDM, BD3LM, and E2D2.

| family | model | trainer | evaluator | |---|---|---|---| | ARM | lit_gpt/model.py (GPT) | pretrain/train_ar.py | eval/evaluate_ar.py | | MDM | lit_gpt/diffmodel.py (TransEncoder) | pretrain/train_mdm.py | eval/evaluate_mdm.py | | BD3LM | lit_gpt/bd3lm_model.py (BD3LM) | pretrain/train_bd3lm.py | eval/evaluate_bd3lm.py | | E2D2 | lit_gpt/e2d2_model.py (E2D2Model) | pretrain/train_e2d2.py | eval/evaluate_e2d2.py |

Environment Setup

Every model in the paper was trained and evaluated with PyTorch 2.4.1 on CUDA 12.1, and the code additionally needs Flash Attention 2.6.3 together with its rotary, layer_norm and xentropy CUDA extensions, xFormers 0.0.28, and lm-evaluation-harness 0.4.4. The [Dockerfile](Dockerfile) is the reference definition of that environment, and building it is the one path that needs no further steps:

docker build -t dlm-scaling .

[environment.yaml](environment.yaml) is the fully resolved lock of the same environment, down to transitive dependencies:

conda env create -f environment.yaml

It installs everything except Flash Attention and its three CUDA extensions, which are not available from PyPI and have to be compiled locally; step 3 of [CONDA.md](CONDA.md) does that, and the rest of that file walks through the same environment one package at a time, at the same pins, if you would rather not use either of the two files above.

Data

All models are trained on SlimPajama. The preprocessing script streams the dataset from HuggingFace, tokenizes it with the TinyLlama sentencepiece tokenizer, and packs the result into binary chunks of 2049 × 1024 tokens:

python scripts/prepare_slimpajama.py --split train --destination data/slim_star_combined
python scripts/prepare_slimpajama.py --split validation --destination data/slim_star_combined

The trainers look for those chunks in data/slim_star_combined unless DATA_DIR points somewhere else. Following TinyLlama, the GitHub and code subset is dropped from every split.

Pretraining

A run is fully specified by the model size it trains (--model) and the compute budget it spends (--flops, in units of 10¹⁸ FLOPs). The budget is counted in real training FLOPs, using the per-variant formulas in lit_gpt/flops.py rather than the 6*ND* approximation, so the same --flops value means the same amount of compute across all four families.

Everything else is held fixed across the study. Each run occupies one node of 8 GPUs and sees a global batch of 256 × 2048 tokens. The optimizer is AdamW with β = 0.9/0.95, a weight decay of 0.1 and gradients clipped at 1.0, and the learning rate warms up linearly over 1% of the steps (at least 100 of them) to a peak of 4e-4 before decaying on a cosine schedule to a tenth of that peak. Every diffusion variant is trained with the low-variance loss, which --low_var_loss selects.

# ARM, 944M parameters at C = 1e20 FLOPs
lightning run model --node-rank=0 --accelerator=cuda --devices=8 --num-nodes=1 \
pretrain/train_ar.py --model 944 --flops 100.

# MDM, 944M at C = 1e20
lightning run model --node-rank=0 --accelerator=cuda --devices=8 --num-nodes=1 \
pretrain/train_mdm.py --model 944 --flops 100. --low_var_loss

# BD3LM, 142M at C = 1e19, block size b = 4
lightning run model --node-rank=0 --accelerator=cuda --devices=8 --num-nodes=1 \
pretrain/train_bd3lm.py --model 142 --flops 10. --mtp_window 4 --low_var_loss

# E2D2, 472M at C = 1e20, block size b = 32, 14 decoder layers
lightning run model --node-rank=0 --accelerator=cuda --devices=8 --num-nodes=1 \
pretrain/train_e2d2.py --model 472 --flops 100. --mtp_window 32 --n_dec_layers 14 \
--low_var_loss

When training finishes, each script scores the validation set and logs the readouts that the scaling laws are fitted on. metric/val_loss is the exact NLL for the ARM and the ELBO for the diffusion variants, and the diffusion variants additionally log metric/duel_nll, the exact likelihood of the validation data under the...