LLaDA-8B-MDLM-LoRA

LoRA adapter for GSAI-ML/LLaDA-8B-Base trained for 2,000 steps on Nemotron-SFT-Code with the MDLM objective. It is the absorbing-only control arm of the 8B transfer experiment of Corrective Diffusion Language Models (NeurIPS 2026): a matched, compute-equal comparison of MDLM and CDLM fine-tuning in which the two arms share the corpus, steps, schedule, seed, batch and LoRA configuration and differ only in mixture_prob and noise_token_wt. The other arm is Shuibai12138/LLaDA-8B-CDLM-LoRA. The experiment will appear in the camera-ready version of the paper; its numbers were reported in the public OpenReview discussion.

step200/ holds the early checkpoint: a separate 200-step run with the same configuration (the same schedule, so its first 200 steps are the same as this run's up to GPU nondeterminism).

Paper: arXiv:2512.15596. Code: zhangshuibai/CDLM.

Training

Base model	GSAI-ML/LLaDA-8B-Base, revision 0f2787f2d87e (8.02B parameters)
Objective	mixture_prob=0, noise_token_wt=0, clean_token_wt=0 (MDLM: absorbing corruption only). The absorbing loss is LLaDA's: t ~ U(0,1) per sequence, mask probability (1-1e-3)t + 1e-3, cross-entropy on masked tokens divided by the mask probability
LoRA	r = 64, alpha = 128, dropout 0.05 on q_proj, k_proj, v_proj, attn_out, ff_proj, up_proj, ff_out inside transformer.blocks.*; 167.77M trainable parameters (2.05%); embeddings and output layer frozen
Data	nvidia/Nemotron-Pretraining-SFT-v1, config Nemotron-SFT-Code, all 78 parquet shards (Hub revision af7991c59eeb, the revision training/data_prep/prepare_nemotron_sft_code.py pins; the shards are shuffled with seed 42 and assigned round-robin to the 4 ranks, and each document is tokenised, given an EOS token and split into chunks of at most 4,096 tokens (chunks under 16 tokens dropped). A 2,000-step run reads the first 3,000 documents of one shard per rank: 12,000 documents, 3.55M tokens, 296 tokens per document on average
Steps	2,000
Optimiser	AdamW, betas (0.9, 0.95), weight decay 0.01, grad clip 1.0; peak lr 3e-4, 200 warmup steps, cosine over a 20,000-step horizon (lr at step 2,000 = 2.94e-4)
Batch	global 12 sequences of at most 4,096 tokens (micro 1 x grad-accum 3 x 4 GPUs); documents are not packed
Precision	bf16 base weights, fp32 LoRA parameters and optimiser states, activation checkpointing
Seed	42
Hardware	4 x A100-PCIE-40GB, 46.5 min (200-step run: 4.7 min)
Code	training/llada8b_lora/run_train.sh mdlm 2000 full; the trainer file used on 2026-07-25 is byte-identical to train_llada_mixture.py at tag v1.0.2-corrective-training; v1.0.4-corrective-training adds only the pinned base-model revision, which is the one this run loaded. The 200-step run is equivalent to run_train.sh mdlm 200 gate; it was launched from an earlier copy of the script that passed --save_steps 1000 instead of --save_steps 0, which is why step200/train_config.json records save_steps: 1000. No checkpoint falls within 200 steps, so this has no effect on the adapter

The trainer's DataLoader runs two workers over an unsharded IterableDataset, so each chunk is used twice in a row with independent noise; 2,000 steps see 12,000 distinct chunks. This is how every LLaDA-8B adapter of the release was trained. The loss is computed with PyTorch cross-entropy, so the liger-kernel gradient issue of the 0.5B trainer does not apply. Training is not bitwise deterministic on GPUs: a re-run matches the step-1 loss exactly, but by step 20 its loss differs by 5e-4 to 7e-3 relative in the re-runs we compared, and the difference grows after that (the Nemotron CDLM 200-step run differs from the 2,000-step run by 0.7% at step 20 and 1.9% at step 200). Use this adapter, not a retrained one, to reproduce the numbers exactly.

train_config.json and train_log.json are the configuration and loss log the trainer wrote for this run.

Results

Localisation is measured on CRB HumanEval with n_replace=1 and all three error types (n = 541). Confidence is the probability the model assigns to the token currently at each position, which is the quantity the self_conf-remask:vanilla sampler uses; the gap is mean confidence on clean positions minus mean confidence on the injected error positions, and Top-K is the fraction of samples with an error among the K least confident positions. Pass@1 (%) is CRB repair at T=1 and T=4 refinement steps with confidence threshold 0.9, remove_all, temperature 0 and batch size 1, macro-averaged over {HumanEval, HumanEval+, MBPP, MBPP+} x {operator, var, literal} at n_replace=1.

Raw vs. de-fenced. Fine-tuning on this corpus teaches the model markdown code fences. The HumanEval+ buggy bodies end with a spare blank line, which both fine-tuned models fill with a fence under remove_all refinement (78.7% (MDLM) and 94.9% (CDLM) of HumanEval+ completions at T=1), and the program then fails to parse. defence.py in the code repository truncates each completion at the first fence and strips trailing whitespace, with the same rule for every model; both numbers are given. The rule does not recover CDLM's HumanEval+ cells at T=4 (35.3-38.8 de-fenced, against 71.2-77.7 at T=1): in 221 of its 511 HumanEval+ completions there, the first fence opens right after the docstring, so the text kept before it has no function body. Its de-fenced T=4 mean includes these failures.

model	confidence gap	Top-1 hit (%)	Top-3 hit (%)	Pass@1 T=1, de-fenced	Pass@1 T=1, raw	Pass@1 T=4, de-fenced	Pass@1 T=4, raw
LLaDA-8B-Base (no fine-tuning)	0.265	16.5	54.7	55.1	55.1	59.5	59.5
+ MDLM LoRA	0.402	15.2	73.4	60.0	47.5	60.9	46.0
+ CDLM LoRA	0.777	60.8	92.8	72.5	53.0	62.3	51.8

Early checkpoint (step200/), localisation only:

model (200 steps)	confidence gap	Top-1 hit (%)	Top-3 hit (%)	Top-5 hit (%)
MDLM	0.258	10.5	55.6	80.6
CDLM	0.790	53.0	87.4	97.0

Pass@1 differences below about 0.5 points are within the evaluation noise (execution timeouts).

Code-generation guard (training/llada8b_lora/run_codegen_guard.sh): from-scratch generation on the 164 HumanEval+ and 378 MBPP+ problems of the lm-eval tasks used for the 0.5B models, one greedy sample per problem (128 tokens, 128 steps, the repository sampler), pass@1 in %. MBPP+ is scored on the three original asserts the prompt shows, as in the 0.5B evaluation, and with the EvalPlus plus tests.

model	HumanEval+	MBPP+	MBPP+ plus tests
LLaDA-8B-Base (no fine-tuning)	28.0	52.9	44.7
+ MDLM LoRA	7.3	27.0	19.6
+ CDLM LoRA	7.3	24.6	19.6

Post-training on either corpus lowers these scores for both objectives; the corrective and the absorbing-only adapter trained on the same data do not differ significantly (exact McNemar tests in the code repository's README).

Usage

The adapter is a PEFT LoRA adapter for the bidirectional masked diffusion model LLaDA-8B-Base. Merge it into the base model before use:

import torch
from transformers import AutoModel, AutoTokenizer
from peft import PeftModel

base = AutoModel.from_pretrained("GSAI-ML/LLaDA-8B-Base", revision="0f2787f2d87eac5eed8a087d5ecd24277e6255b2",
                                 trust_remote_code=True, dtype=torch.bfloat16)
tok = AutoTokenizer.from_pretrained("GSAI-ML/LLaDA-8B-Base", revision="0f2787f2d87eac5eed8a087d5ecd24277e6255b2", trust_remote_code=True)
model = PeftModel.from_pretrained(base, "Shuibai12138/LLaDA-8B-MDLM-LoRA").merge_and_unload().eval()

To reproduce the numbers above, use the code repository (zhangshuibai/CDLM, training/llada8b_lora/); its scripts accept a Hub adapter id, optionally with a subfolder and a revision (<owner>/<name>[/<subfolder>][@<revision>]). The commands below pin the revision that holds these weights.

git clone --branch v1.0.4-corrective-training https://github.com/zhangshuibai/CDLM && cd CDLM
pip install -r training/llada8b_lora/requirements.txt
bash training/llada8b_lora/fetch_crb_inputs.sh          # the 24 CRB input files of this experiment
# localisation (confidence gap and Top-K hit rate)
python training/llada8b_lora/eval_confidence.py --label mdlm --adapter Shuibai12138/LLaDA-8B-MDLM-LoRA@d9f1ca8f75b9ce7437781376cab07640d459db67 --out outputs/conf_mdlm.json
# one CRB repair cell (HumanEval, operator errors, n_replace 1, T=1)
bash training/llada8b_lora/run_crb_cell.sh 0 mdlm Shuibai12138/LLaDA-8B-MDLM-LoRA@d9f1ca8f75b9ce7437781376cab07640d459db67 human-eval operator 1 2
# all cells of the base model and both arms, then the Pass@1 table
MDLM_ADAPTER=Shuibai12138/LLaDA-8B-MDLM-LoRA@d9f1ca8f75b9ce7437781376cab07640d459db67 \
CDLM_ADAPTER=Shuibai12138/LLaDA-8B-CDLM-LoRA@eea470e0d03e08030fdbdac4c6944431f8ec0c87 \
GPUS="0 1 2 3" bash training/llada8b_lora/run_crb_parallel.sh 1
python training/llada8b_lora/aggregate_crb.py --n_replace 1

For the 200-step adapter use Shuibai12138/LLaDA-8B-MDLM-LoRA/step200@d9f1ca8f75b9ce7437781376cab07640d459db67 in the scripts, or subfolder="step200" in PeftModel.from_pretrained.

Licence and attribution

Adapter: MIT. Base model: GSAI-ML/LLaDA-8B-Base (MIT). Training data: nvidia/Nemotron-Pretraining-SFT-v1, which is subject to its own licence terms. An adapter trained with the same code on the public nvidia/OpenCodeInstruct corpus is Shuibai12138/LLaDA-8B-MDLM-LoRA-OpenCodeInstruct.

Citation

@inproceedings{zhang2026corrective,
  title     = {Corrective Diffusion Language Models},
  author    = {Zhang, Shuibai and Peng, Fred Zhangzhi and Zhang, Yiheng and Pan, Jin and Chrysos, Grigorios G.},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026}
}