plaidq

<h1 align="center">PlaidQ</h1>

Distilled Continuous Diffusion Language Models Can Write Code in Few Steps—or One

PlaidQ is a 0.7B continuous diffusion language model that writes working code in a
few denoising steps &mdash; or in one. No token-by-token decoding: the whole program is denoised
at once, and distillation compresses the trajectory down to a single global update.

https://github.com/user-attachments/assets/09377dd5-1a48-48ec-97a2-61915bb0937a

Few-step Code Generation Demo.

Quickstart

A ready-made environment is in the Dockerfile.

git clone https://github.com/pengzhangzhi/plaidq && cd plaidq
pip install -e .

python -m plaidq.sample \
    --checkpoint_path=hf://fredzzp/plaidq-0.7b-16step \
    --prompt="def quicksort(arr):" \
    --sampling_timesteps=16 --score_temp=0.5 \
    --n_samples=4

For the one-step model:

python -m plaidq.sample \
    --checkpoint_path=hf://fredzzp/plaidq-0.7b-1step \
    --prompt="def quicksort(arr):" \
    --sampling_timesteps=1 --score_temp=0.99945 \
    --n_samples=4

Pass --score_temp with the step budget. It is not optional tuning: the score temperature
has to rise as the budget falls, and the flag's default is a single fixed 0.6 that is tuned for
no particular budget. At 16 steps 0.6 is close enough to the tuned 0.5 to look fine; at one step
it does not produce code at all &mdash; it produces a repeated token. The tuned values live in
configs/steps/base_s<N>.json; the two above are the ones those files record for 16 and 1. One
step is also the hard operating point on purpose: expect a plausible program more often than a
correct one, and read the numbers off the paper's pass@10 rather than off a single sample.

--checkpoint_path also takes a local .pt file or a directory exported by
scripts/export_to_hf.py, so a model you train yourself by following
Train, distill, evaluate drops straight in.

Useful flags: --sampling_timesteps (denoising steps), --response_tokens (generation budget,
default 128), --score_temp, --sampler=ddim|ancestral, --seed.

Checkpoints

All three are on the Hub and are downloaded on first use &mdash; pass the hf:// spec to any CLI
here that takes a checkpoint.

Train, distill, evaluate

Get the data first (details in scripts/README.md; corpora land under
$PLAIDQ_DATA_ROOT, default ./data, benchmarks under $PLAIDQ_EVAL_DATA, default
./data/eval).

You need a Hugging Face token, and access granted in advance. Request it on each dataset's
page, then huggingface-cli login (or export HF_TOKEN). The code-only path needs two
approvals &mdash; nvidia/Nemotron-Pretraining-Code-v1 and nvidia/Nemotron-Pretraining-SFT-v1
&mdash; and both are granted manually, so ask before you need them. The full blend adds
Nemotron-CC-v2.1 (manual) and Nemotron-CC-Math-v1 (automatic);
Nemotron-Pretraining-Specialized-v1 and the benchmark data are ungated. File listings on these
repos are public while their contents are not, so a missing approval only shows up once the
download starts.

Set PLAIDQ_DATA_ROOT before you start. It defaults to ./data, i.e. inside the clone,
and the download below is hundreds of gigabytes &mdash; point it at a volume that can hold it.

Start with the code-only corpus. It is 262 GiB rather than 1.5 TiB, and it is a recipe we
have actually trained on, not a suggestion &mdash; it produces a code specialist rather than a
reproduction of the released teacher, which is what most people building on this want anyway.

export PLAIDQ_DATA_ROOT=/somewhere/with/room     # NOT the default ./data inside the repo

python scripts/download_nemotron_mix.py --preset=code   # -> $PLAIDQ_DATA_ROOT/nemotron_mix/raw
python scripts/download_nemotron.py     --preset=code   # -> $PLAIDQ_DATA_ROOT/nemotron/raw
python scripts/download_eval_tasks.py                   # HumanEval / MBPP / infilling splits

python scripts/build_mix_index.py \
    --mix_root=$PLAIDQ_DATA_ROOT/nemotron_mix --sft_root=$PLAIDQ_DATA_ROOT/nemotron
export PLAIDQ_MIX_MANIFEST=$PLAIDQ_DATA_ROOT/nemotron_mix/meta.json

Drop --preset=code from the two corpus downloads to get the full blend the released teacher
used instead. Nothing else changes: build_mix_index.py indexes whatever is on disk and skips
what is not, so the same command builds either manifest and prints which domains it found.

The released recipe

Three runs, in order. Each one reads the checkpoint the one before it wrote.

The distillations are short (4k steps) because they start from a trained teacher. Almost all the compute is
in the pretraining. Both distill commands below leave --steps at its default of 20,000; the
released students were taken at early steps such as 4k steps.

Every flag below is one the released checkpoints actually record, and several are not the
repo defaults &mdash; copy them as written. Anything not listed is a default and can be left
alone. The torchrun lines launch one 8-GPU node, which is the shape both distillations
actually used; only pretraining needs more, and multi-node launch is up to your scheduler.

1. Pretrain the teacher.

torchrun --nproc_per_node=8 train.py \
    --run_dir=./runs/plaidq \
    --seq_len=2048 --batch_size=8 --steps=380000 \
    --prefix_conditioning=True --prefix_boundary_prob=0.5 \
    --weight_decay=4e-5 --val_batch_size=8 --auto_resume=True

--batch_size is the per-GPU micro-batch. Global batch = batch_size x grad_accum_steps x
world size, so 8 x 1 x 64 = 512, which is the released teacher's
global batch; on fewer GPUs raise --grad_accum_steps to keep that product at 512.

Three flags that are easy to miss:

• --prefix_boundary_prob=0.5 and --weight_decay=4e-5 change the trained model, and neither is
  a default.
• --val_batch_size defaults to 64 &mdash; a validation batch of 64 x 2048 tokens, which runs out
  of memory on an 80 GB card. The failure lands at the first validation, not at startup.
• --auto_resume=True is what writes checkpoint_latest.pt, the file the next two runs read.

The rest is default: lr 1.4e-3, warmup to step 2500 then decay to zero, embed_dim=16,
self-conditioning on, no EMA.

2. Distill to few steps (distribution matching). Student and critic both start from the
teacher. Both learning rates are 30x below the defaults, and the rollout mixture puts all its
weight on length 4; both matter.

torchrun --nproc_per_node=8 train_distill.py \
    --teacher_path=./runs/plaidq/checkpoint_latest.pt \
    --run_dir=./runs/plaidq-16step \
    --lr=1e-6 --fake_score_lr=1e-6 --ema=0 \
    --rollout_len_weights='[0,0,0,1,0,0,0,0]' \
    --milestone_every=500

--milestone_every=500 is what produces milestone_3999.pt, the released few-step student.

3. Distill to one step (paired-trajectory regression). This run needs a rollout teacher &mdash;
a frozen copy of the few-step student that generates the paired targets. Unwrap it from a
milestone first, because DistillConfig rejects a --reg_target_path still named milestone_*.

python scripts/extract_ema_student.py --allow_online \
    --milestone=./runs/plaidq-16step/milestone_3999.pt \
    --out=./runs/plaidq-16step/rollout_teacher.pt

torchrun --nproc_per_node=8 train_distill.py \
    --teacher_path=./runs/plaidq/checkpoint_latest.pt \
    --reg_target_path=./runs/plaidq-16step/rollout_teacher.pt \
    --run_dir=./runs/plaidq-1step \
    --lr=1e-4 --fake_score_lr=1e-4 --ema=0 \
    --batch_size=4 --grad_accum_steps=8 --student_update_freq=1 \
    --dmd_loss_weight=0 --reg_kind=ce --reg_weight=1.0 \
    --rollout_len_weights='[1,0,0,0,0,0,0,0]' \
    --token_metric_rollout_steps=1 --milestone_every=500

--allow_online is required: run 2 sets --ema=0, so its milestones hold no EMA shadow and the
online student is the thing to extract. Drop the flag if you trained with an EMA.
--dmd_loss_weight=0 switches the DMD term off and leaves regression alone, so --reg_weight has
nothing to trade against and 1.0 is simply what the released run used.

Both distillations ran on one 8-GPU node, for a
global batch of 256 each (8 x 4 x 8 in run 2, 4 x 8 x 8 in run 3). Only the pretraining run
needs more than a single node.

Evaluate. configs/paper/full4_s128.json is the reported protocol: all four code
benchmarks, zero-shot, 20 samples per problem, pass@10, a 128-token response budget, and a
first-order DDIM sampler at score_temp=0.5 on a gamma-uniform grid. Pass it explicitly —
running with no --config_json falls back to the built-in defaults in plaidq/eval/config.py,
which are a different operating point (two tasks, 512 steps, score_temp=0.6) and will not
reproduce the tables. Prefix with torchrun to shard it across GPUs.

python -m plaidq.eval --ckpt=./runs/plaidq/checkpoint_latest.pt \
    --config_json=configs/paper/full4_s128.json --out=results.json

For a different step budget, sampler or task list, pass another --config_json.
configs/steps/ sweeps the budget from base_s1 to base_s512 on a two-task subset
(HumanEval and MBPP+) rather than all four benchmarks, so those numbers are not directly
comparable to the table above; copy one and edit its tasks or steps to change that. Note
that score_temp is budget-dependent and each config carries its own &mdash; the value that is
right at 16 steps is not right at 1.

Evaluation runs model-generated code. See docs/EVAL.md before pointing it at
a machine you care about.

Citation

If you use PlaidQ in your work, please cite:

@misc{peng2026distilledcontinuousdiffusionlanguagemodels,
  title         = {Distilled Continuous Diffusion Language Models Can Write Code in Few Steps---or One},
  author        = {Fred Zhangzhi Peng and Kaiwen Zheng and Anru R. Zhang},
  year          = {2026},
  eprint        = {2609.04531},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2609.04531}
}