DiffuRefill-1B-base / paper /diffurefill.typ
Asilarkness's picture
DiffuRefill-1B base: weights, code, evaluation, technical report
9a57d8a verified
Raw History Blame Contribute Delete
20.9 kB
#set document(title: "DiffuRefill-1B: Pretraining a Masked Diffusion Language Model on Free, Preemptible Notebook GPUs", author: "Asilarkness")
#set page(paper: "a4", margin: (x: 2.1cm, y: 2.2cm), numbering: "1")
#set text(font: ("New Computer Modern", "Libertinus Serif"), size: 10pt)
#set par(justify: true, leading: 0.62em)
#set heading(numbering: "1.1")
#show heading.where(level: 1): it => block(above: 1.3em, below: 0.7em, text(size: 12pt, weight: "bold", it))
#show heading.where(level: 2): it => block(above: 1.0em, below: 0.5em, text(size: 10.5pt, weight: "bold", it))
#show table: set text(size: 8.6pt)
#show figure.caption: set text(size: 8.8pt)
#align(center)[
#text(size: 15pt, weight: "bold")[DiffuRefill-1B: Pretraining a Masked Diffusion\ Language Model on Free, Preemptible Notebook GPUs]
#v(0.6em)
#text(size: 11pt)[Asilarkness]
#v(0.2em)
#text(size: 9pt)[Technical report · October 2026 · #link("https://huggingface.co/Asilarkness/DiffuRefill-1B-base")[huggingface.co/Asilarkness/DiffuRefill-1B-base]]
]
#v(0.8em)
#block(inset: (x: 1.2cm))[
#text(weight: "bold")[Abstract.]
We pretrain DiffuRefill-1B, a 1.08B-parameter masked-diffusion language model, for 200,000 steps on 92.8B tokens without paying for any compute. Training ran for 21 days on free notebook GPUs that are reclaimed without warning, often several at once, and typically live for hours. The only state shared between machines is a public model repository on the Hugging Face Hub: workers take 150 local AdamW steps, push their weights, and a merger that can itself be killed and restarted anywhere averages them into the next global model (DiLoCo with plain replacement averaging). We describe the design that made this survive about 200 worker launches and repeated loss of the whole fleet, and we report the measurements that shaped the run, including the negative ones: keeping local drift at merge time, outer Nesterov momentum, a mid-run switch to Muon, span and PMI masking, phrase tokens and cautious AdamW all failed to help, while halving the learning rate for the merge interval, document masking, short rows and weight averaging did. On standard zero-shot benchmarks the model is weaker than autoregressive models trained with 3-280x more compute (HellaSwag 34.8, ARC-Easy 52.4, PIQA 62.7, SciQ 89.5), which we report as it is. As a diffusion model it generates 1.9x faster than comparable autoregressive models at batch size 1 when committing one token per pass, and up to 15x faster when committing eight. Weights, code and evaluation scripts are released.
]
= Introduction
Pretraining even a small language model is usually taken to require a budget: thousands of GPU-hours on machines that stay up. This report asks what a single person can do with none. The resource we had was free notebook GPUs (RTX PRO 6000 Blackwell, 96 GB, on the molab service), which come with two constraints that rule out ordinary data-parallel training: a box can be reclaimed at any moment, without notice, and boxes cannot talk to each other. In practice most boxes lived from one to a few hours, whole fleets of eight disappeared within the same minute several times, and new boxes arrived only when a person opened new notebooks.
We therefore built training around the one thing that persists: a public model repository on the Hugging Face Hub. Each worker trains on its own for a fixed number of steps, uploads its weights, and downloads the next average. This is DiLoCo @diloco with the Hub as parameter server and with every component, including the merger, allowed to die.
We chose a *masked diffusion* language model (MDLM @mdlm, @smdm) rather than an autoregressive one. Diffusion models predict any subset of positions from bidirectional context, can commit several tokens per forward pass, and are a natural fit for our longer-term goal: a model that first writes out its reasoning and then fills in an answer in parallel (§8). The price, known from prior work @smdm, is lower compute efficiency on likelihood-style benchmarks.
Our contributions are:
- *A fault-tolerant training recipe* for preemptible, mutually unreachable GPUs, using only a model hub for coordination (§3).
- *DiffuRefill-1B*, a 1.08B masked diffusion LM trained to 200k steps on 92.8B tokens at zero monetary cost, released with code (§4).
- *A record of measured interventions*, positive and negative, at 38M-parameter "stand" scale and at 1B (§5).
- *An honest evaluation* against autoregressive models on the same harness, and a generation-speed comparison (§6).
= Model
DiffuRefill-1B is a pre-norm transformer with bidirectional attention (Table 1). It is trained with the MDLM objective @mdlm: for each row a mask rate $t tilde U(0, 1)$ is drawn, each token is replaced by `[MASK]` with probability $t$, and the model is trained with cross entropy on the masked positions.
#figure(
table(columns: 2, align: (left, left), stroke: 0.4pt,
[Parameters], [1.08B (embedding tied to the output head)],
[Layers / width / heads], [18 / 2048 / 16 (head dim 128)],
[FFN], [SwiGLU, 5632],
[Normalisation, positions], [RMSNorm, RoPE (base $10^4$)],
[Context], [2048 tokens],
[Tokenizer], [MiniCPM4 (73,440 pieces; model vocabulary 73,760 incl. `[MASK]`)],
[Objective], [MDLM masked diffusion, independent token masking],
),
caption: [Architecture.],
)
Three data-side choices changed the objective in ways we measured (§5):
*Document masking.* Rows are packed documents, and with bidirectional attention a masked position can be filled from an unrelated neighbouring document. We restrict attention to the token's own document with FlexAttention block masks, which was also faster than full attention (13.2 ms vs 17.5 ms per call at the training shape) because blocks between documents are skipped.
*Salient span masking.* On 25% of rows, capitalised words that do not start a sentence, and numbers, are masked as whole spans, so that a name must be recalled from the sentence rather than completed from its visible half.
*Short rows.* 10% of micro-batches are reshaped into rows of 128 tokens, so the model also sees short contexts that begin at a document start.
= Training system
#figure(
table(columns: 2, align: (left, left), stroke: 0.4pt,
[Worker], [loads the latest global, trains $H=150$ local AdamW steps (batch 64 × 2048 tokens, 131k tokens per step), uploads `ckpt_wNN.pt` with its round tag and a small log],
[Merger], [waits until at least $N$ fresh checkpoints carry the current round tag (normally $N=2$; $N=1$ when one worker is left), allows a 150 s grace for stragglers, averages them, publishes `global_b.pt` and a new round tag],
[Outer step], [plain replacement: the new global is the mean of the workers (outer learning rate 1, no momentum)],
[Side averages], [the merger also keeps an EMA ($beta = 0.9$) of globals and, in the final stretch, a uniform average (SWA) and periodic snapshots],
[Recovery], [a lost merger is restarted on any surviving box; it first merges leftover checkpoints of dead workers ($N=1$), then waits for new workers; workers seed from the last global],
),
caption: [The hub-synchronised DiLoCo loop. The Hub repository is the only shared state.],
)
A round takes about 9.5 minutes regardless of how many workers take part, because every worker runs the same 150 steps; more workers add tokens per round, not rounds per hour. About 26% of wall time was spent waiting at synchronisation. A worker processes about 35k tokens/s on one RTX PRO 6000 with `torch.compile` per block (17% faster than eager; FP8 via torchao gave no further gain on this GPU).
*Failure handling.* Three rules carried the run through about 200 worker launches. (i) Liveness is read from the logs workers push, not from whether a notebook answers: endpoints drop connections in bursts, and a reachability check twice declared a healthy fleet dead. (ii) After a total fleet loss, the merger is started before any worker is seeded, so that checkpoints already uploaded by dead workers are merged rather than overwritten. (iii) A rollback is published under a new round number: a stale worker checkpoint carrying the same round tag as a re-published global was once merged into it. The repository history is periodically squashed to stay within storage limits.
= Data and schedule
For the first 158.8k steps the data was Ultra-FineWeb-L3 @ultrafineweb (English, the Multi-Style and QA synthetic rewrites of filtered web text), later with 5% of fact sentences generated by an open model. For the final 41k steps we annealed on a higher-quality mix (Table 3), streamed directly from the Hub with a custom reader that keeps memory bounded (the largest sources store up to 2.5 GB per decoded Parquet row group, which killed the first attempt).
#figure(
table(columns: 3, align: (left, right, left), stroke: 0.4pt,
[*Source*], [*Share*], [*Content*],
[Ultra-FineWeb-L3 (continuation)], [18%], [rewritten web text],
[Dolmino mix (OLMo 3 stage 2) @olmo3], [20%], [math, code, QA, science, high-quality web],
[Nemotron Pretraining Specialized: Fact-Seeking], [15%], [fact Q&A from Wikipedia],
[Nemotron Pretraining Specialized: Wiki-Rewrite], [15%], [rewritten Wikipedia],
[Nemotron-CC-Math (4+)], [8%], [math web pages],
[UltraData-Code L3], [7%], [code with explanations],
[OpenCoder annealing corpus], [8%], [algorithmic corpus, synthetic snippets],
[UltraData-Math textbook exercises], [4%], [synthetic exercises],
[Generated fact sentences], [5%], [our own, open-model generated],
),
caption: [Anneal mix for steps 158.8k-200k.],
)
The learning rate warmed up over 500 steps to $3 dot 10^(-4)$ and followed a cosine to 10% over 200k steps. It was halved to $1.5 dot 10^(-4)$ early in the run (§5.1). From step 180k it decayed linearly to zero at 200k @d2z, and the released weights are the uniform average of the 130 merged globals from step $approx$181k to 200k. AdamW used $beta = (0.9, 0.95)$, weight decay 0.1 and gradient clipping at 1.0. In total the model saw 92.8B tokens (summed over workers), about $6.0 dot 10^20$ training FLOPs by the $6 N D$ rule.
= What we measured
All 1B measurements below use a fixed held-out batch scored with fixed masks at rates 0.3/0.6/0.9, so the difference between two checkpoints carries no sampling noise. Training loss alone was misleading more than once: it stayed flat while held-out quality got worse (§5.1). Stand experiments use a 38M-parameter model with identical seeds per arm.
== Distributed optimisation
*Replacement beats keeping local drift.* We tried adopting the global as $theta arrow.l theta + (g - theta_"pushed")$, which keeps each worker's progress since its last push. Adoption then made the model worse in 60 of 63 probes (mean +0.079 CE) and the merged model drifted from 4.41 to 4.60 over 16 rounds; plain replacement improved it by 0.08-0.10 per round. Averaging works because the components the workers disagree on are divided by $N$; keeping local drift puts them back.
*The learning rate must fit the merge interval.* At $3 dot 10^(-4)$ the merged model sat at 4.33 for about 1.3B tokens. The probe showed why: 150 local steps raised held-out CE by 0.076 and the merge exactly undid it. Halving the rate cut local drift to 0.013 and the plateau ended.
*Outer Nesterov hurts here.* With momentum 0.5 (outer step 1.6x plain averaging), the global rose three rounds running, 2.9657 to 2.9721, and stopped rising the round it was switched off.
*EMA helps for evaluation, not for training.* An EMA of globals was at or below the global at every round once it held 4+ globals (2.9533 vs 2.9555 at $n=20$). Re-seeding the fleet from it gave the gain back within a round.
== Objective and data
#figure(
table(columns: 3, align: (left, left, left), stroke: 0.4pt,
[*Intervention*], [*Scale*], [*Result*],
[Span masking on half of the batches], [38M], [worse than independent masking: CE 3.748 vs 3.651; adopted earlier on weak evidence, then removed],
[PMI (collocation) masking], [38M], [worse than independent masking at every coverage],
[Phrase tokens (PMI word pairs merged into the vocabulary)], [30M], [faster early, worse at equal FLOPs: 1.891 vs 1.860 bits/byte],
[Short rows (10% of micro-batches at 128 tokens)], [1B], [better on long CE (-0.004), short-context CE (-0.021) and fact probe; adopted],
[Document masking], [1B], [adopted; also faster],
[Self-conditioning], [38M / 1B], [-0.038 / -0.008 CE, but +27% step cost and no fact gain; not adopted],
[N-gram (Engram) memory], [1B mid-run], [early fact edge faded within 1000 steps; CE slightly worse],
[Product-key memory (v2, query BN, slow keys)], [1B mid-run], [fact probe slightly better at every eval, CE flat; not adopted],
),
caption: [Objective, tokenisation and architecture interventions.],
)
== Optimisers
*Muon* @muon beat AdamW from scratch on the stand by a wide margin (held-out CE 3.748 to 3.177, reaching AdamW's final loss in under half the steps). Switched onto the 1B at step 84k, it raised one worker's loss monotonically for a thousand steps (3.16 to 3.32) while an AdamW worker beside it held 3.13-3.15; re-tested on the stand with RMS-matched updates @moonlight, the mid-run switch gained nothing (+0.0085). *Cautious AdamW* @cautious improved the stand by 0.197 CE but changed the 1B by $+0.0001 plus.minus 0.0012$. *Outer orthogonalisation* (replacing the averaged update by its polar factor) helped on a 4-worker stand run (held-out CE 3.750 vs 4.005 for plain averaging after 60 rounds, against 3.407 for synchronous training with the same total batch), but was not tried on the 1B and needs more seeds.
== The final stretch
Linear decay to zero from step 180k and the four weight averages we kept were compared at step 200k on the same probes (Table 5). All five are within 0.0004 CE; the uniform average is best on every probe and is released. We also note that the EMA at step 183k already scored 2.9495, i.e. the last 17k steps did not move these probes, even though the learning rate went to zero.
#figure(
table(columns: 7, align: (left, right, right, right, right, right, right), stroke: 0.4pt,
[*Weights*], [*held-out*], [*wiki*], [*code*], [*fact mc*], [*fact nll*], [*fact EM*],
[global (step 200k)], [2.9503], [3.5397], [1.6668], [.558], [5.899], [9.6%],
[EMA, $beta=0.9$], [2.9503], [3.5397], [1.6667], [.558], [5.899], [9.6%],
[*uniform average, 130 globals*], [*2.9501*], [*3.5393*], [*1.6639*], [*.561*], [*5.894*], [9.6%],
[average of last 3 snapshots], [2.9502], [3.5398], [1.6656], [.557], [5.900], [9.7%],
[average of last 5 snapshots], [2.9503], [3.5394], [1.6641], [.559], [5.895], [9.5%],
),
caption: [Candidates for the released weights (masked CE; PopQA fact probe of 1050 facts).],
)
= Evaluation
== Zero-shot benchmarks
All models are scored with lm-evaluation-harness @lmeval, zero-shot, on the same tasks. For autoregressive models we use the standard HF backend. Our model is plugged into the same harness: the log-likelihood of a continuation is computed by the left-to-right chain rule, predicting token $i$ with tokens before it visible and tokens from $i$ on masked. On 1000 HellaSwag items this scorer gave 42.5 against 40.8 for the Monte Carlo ELBO estimator of LLaDA @llada (32 samples), so it does not understate the model. TinyLlama and OLMo-1B were also run but load incorrectly under the current `transformers` (e.g. OLMo: LAMBADA 0%) and are omitted.
#figure(
table(columns: 12, align: (left,) + (right,) * 11, stroke: 0.4pt,
[*Model*], [*Params*], [*Tokens*], [*FLOPs*], [*HS*], [*ARC-e*], [*ARC-c*], [*PIQA*], [*WG*], [*OBQA*], [*SciQ*], [*LAMB*],
[*DiffuRefill-1B (ours)*], [1.08B], [0.09T], [$6.0 dot 10^20$], [34.8], [52.4], [25.0], [62.7], [50.6], [27.4], [*89.5*], [34.0],
[Pythia-1B @pythia], [1.0B], [0.3T], [$1.8 dot 10^21$], [47.2], [56.9], [27.0], [69.4], [54.1], [31.2], [83.4], [55.8],
[SmolLM2-360M @smollm2], [0.36B], [4T], [$8.6 dot 10^21$], [56.3], [70.3], [38.1], [71.5], [58.9], [37.8], [91.1], [53.4],
[Qwen2.5-0.5B @qwen25], [0.49B], [18T], [$5.3 dot 10^22$], [52.2], [64.9], [31.8], [69.9], [56.1], [35.0], [93.0], [52.0],
[SmolLM2-1.7B @smollm2], [1.7B], [11T], [$1.1 dot 10^23$], [71.2], [77.8], [47.1], [77.6], [65.9], [44.0], [93.3], [67.5],
[Qwen2.5-1.5B @qwen25], [1.5B], [18T], [$1.7 dot 10^23$], [67.8], [75.5], [45.0], [75.9], [63.5], [40.6], [94.3], [62.2],
),
caption: [Zero-shot accuracy (%). HS, ARC-c, PIQA, OBQA: length-normalised accuracy; others: accuracy. FLOPs by $6 N D$. MMLU, zero-shot: ours 26.9, Pythia-1B 23.1, SmolLM2-360M 25.5, Qwen2.5-0.5B 47.5, SmolLM2-1.7B 48.5, Qwen2.5-1.5B 59.7.],
)
The model is clearly weaker than autoregressive models of similar size, all of which were trained with 3x (Pythia) to 280x (Qwen2.5-1.5B) more compute. Two factors are expected to account for most of the gap: the token budget, and the known compute inefficiency of masked diffusion on likelihood benchmarks (Nie et al. @smdm estimate that a masked diffusion model needs about 16x the compute of an autoregressive one for the same validation loss). For reference, their 1.1B diffusion model trained with $3.3 dot 10^21$ FLOPs (5.5x ours) reports HellaSwag 51.8, ARC-e 48.7, PIQA 69.5 and LAMBADA 52.7, with unsupervised classifier-free guidance at evaluation. Our model is relatively strong on SciQ (above Pythia-1B and SmolLM2-360M), which matches the fact-heavy anneal, and weakest on LAMBADA and WinoGrande; WinoGrande is at chance under both scorers.
== Generation speed
#figure(
table(columns: 4, align: (left, right, right, right), stroke: 0.4pt,
[*Setting (tokens/s, 256 new tokens)*], [*DiffuRefill-1B*], [*SmolLM2-1.7B*], [*Qwen2.5-1.5B*],
[batch 1, one token per pass], [*193*], [101], [82],
[batch 1, 8 tokens per pass], [*1536*], [-], [-],
[batch 32, one token per pass], [431], [*3119*], [2612],
[batch 32, 8 tokens per pass], [*3449*], [-], [-],
),
caption: [Greedy generation on one RTX PRO 6000, bf16. Autoregressive models use HF `generate()` with a KV cache; ours uses no cache and re-reads the 288-token row every pass.],
)
At batch size 1, where latency is set by the number of sequential passes, the diffusion model is about twice as fast even when it commits one token per pass, and 15x faster at eight. At batch 32 the autoregressive models are compute-efficient thanks to the KV cache and our model only catches up at eight tokens per pass. Committing several tokens per pass costs quality: on the stand, committing tokens one at a time in left-to-right order raised distinct-bigram diversity of samples from 0.65 to 0.93 and removed repetition loops, while spacing parallel commits apart cut judge perplexity from 48.0 to 35.1. Optimised AR serving (vLLM and similar) would be considerably faster than HF `generate()`.
= Limitations
This is one run, with interventions decided on small held-out probes and a 38M stand rather than full 1B ablations; several stand wins (Muon, cautious AdamW) did not transfer, and the converse may also hold. The fact probe overlaps with the Wikipedia-derived anneal data. Baselines were evaluated by us with one harness version and default settings and may differ from their published numbers. The model is English-only, is a base model with no instruction tuning or safety training, and generates fluent but often wrong facts.
= Outlook: reasoning
Our goal for the next stage is reasoning. On the stand, a diffusion model that first writes an autoregressive chain of thought and then fills in a computed four-digit answer *in a single parallel pass* went from 45% to 100% accuracy on a multi-step arithmetic task, while the same model without the reasoning trace stayed at 45%. This suggests a natural division of labour for diffusion LMs: sequential reasoning, parallel answering. A recurrent-depth variant (a looped core block) matched an 8-block model's pointer-chasing depth with half the parameters but did not extrapolate with more loops. The next steps are supervised fine-tuning on reasoning traces and reinforcement learning with verifiable rewards.
= Acknowledgements
Compute was provided free of charge by molab notebooks. Storage and coordination ran on the Hugging Face Hub. We thank the authors of Ultra-FineWeb, Nemotron Pretraining Specialized, Dolmino, OpenCoder and UltraData for open data. Engineering, experiments and the drafting of this report were carried out with the help of Claude Code (Anthropic) as a programming assistant.
#bibliography("refs.yml", style: "ieee")