Traj-MC — LLaDA-8B-Base compressed checkpoints

Every compressed LLaDA-8B-Base checkpoint behind the Part-1 table of the Traj-MC study (trajectory-aware calibration for diffusion LLMs), plus the calibration tensors, the evaluation harness, and the results collected so far.

The point of the table is a paired comparison: for each compression method the same algorithm is calibrated twice — once on clean text (clean) and once on masked diffusion-trajectory activations (traj) — and nothing else changes.

Status: three columns are unfinished. svamp, arc_challenge and math500 still need to be run for 22 of the 23 arms. Everything needed to run them is in this repo. See Running the missing columns.

Layout

weights/<family>/<config>/    22 compressed checkpoints (see the table below)
calib/                        calibration tensors (X^T X per target linear) + manifests
code/                         evaluation + compression harness
results/eval_p1/              results already collected (GSM8K, MMLU, BBH, HumanEval, MBPP, …)
results/eval_p1/.queue/       the work queue for the three unfinished columns

Checkpoint format

These are not standalone models. Each directory holds one .pt file per compressed linear, which the harness patches into the dense GSAI-ML/LLaDA-8B-Base at load time. Targets are the 224 linears inside the transformer blocks (32 blocks × {q_proj, k_proj, v_proj, attn_out, ff_proj, up_proj, ff_out}); the top-level unembed head is deliberately excluded.

Family	Files	dtype	Reconstruction
low-rank (base, ours, asvd)	<name>_A.pt, <name>_B.pt	float32	W ≈ A @ B
quant / prune (gptq, awq, sparsegpt, wanda, quant_dllm, sink)	<name>_weight.pt	bfloat16	dense replacement (fake-quant / masked)

<name> is the parameter path with . replaced by _, e.g. model_transformer_blocks_0_attn_out_A.pt. Most directories also carry a compression_summary.json recording the arm, ratio, calibration file and git hash.

The low-rank factors are kept in float32 on purpose — that is the dtype the finished columns were produced with. Do not cast them to bf16 before evaluating or the numbers will not match the published table.

The 23 arms

dense_bf16 is the uncompressed reference and needs no weights.

Arm	Method	Setting	Calib	weights/…
dense_bf16	—	uncompressed bf16	—	(none)
svdllm_m08_c4_clean	SVD-LLM	model ratio 0.8	clean	base/all_m08_c4
svdllm_m08_c4_traj	SVD-LLM	model ratio 0.8	traj	ours/all_m08_c4
svdllm_m06_c4_clean	SVD-LLM	model ratio 0.6	clean	base/all_m06_c4
svdllm_m06_c4_traj	SVD-LLM	model ratio 0.6	traj	ours/all_m06_c4
asvd_m08_c4_clean	ASVD	model ratio 0.8	clean	asvd/all_m08_c4
asvd_m08_c4_traj	ASVD	model ratio 0.8	traj	asvd/all_m08_c4_traj
gptq_w4g128_c4_clean	GPTQ	W4 g128	clean	gptq/w4g128_c4_clean
gptq_w4g128_c4_traj	GPTQ	W4 g128	traj	gptq/w4g128_c4_traj
gptq_w3g128_c4_clean	GPTQ	W3 g128	clean	gptq/w3g128_c4_clean
gptq_w3g128_c4_traj	GPTQ	W3 g128	traj	gptq/w3g128_c4_traj
awq_w4g128_c4_clean	AWQ	W4 g128	clean	awq/w4g128_c4_clean
awq_w4g128_c4_traj	AWQ	W4 g128	traj	awq/w4g128_c4_traj
awq_w3g128_c4_clean	AWQ	W3 g128	clean	awq/w3g128_c4_clean
awq_w3g128_c4_traj	AWQ	W3 g128	traj	awq/w3g128_c4_traj
sparsegpt_unstr50_c4_clean	SparseGPT	50 % unstructured	clean	sparsegpt/unstr50_c4_clean
sparsegpt_unstr50_c4_traj	SparseGPT	50 % unstructured	traj	sparsegpt/unstr50_c4_traj
sparsegpt_unstr70_c4_clean	SparseGPT	70 % unstructured	clean	sparsegpt/unstr70_c4_clean
sparsegpt_unstr70_c4_traj	SparseGPT	70 % unstructured	traj	sparsegpt/unstr70_c4_traj
wanda_unstr50_c4_clean	Wanda	50 % unstructured	clean	wanda/unstr50_c4_clean
wanda_unstr50_c4_traj	Wanda	50 % unstructured	traj	wanda/unstr50_c4_traj
quant_dllm_w2_c4	Quant-dLLM	W2	—	quant_dllm/w2_c4
sink_unstr50	Sink-aware	50 % unstructured	—	sink/unstr50

Calibration

calib/ holds the second-moment statistics the compressors consume, built with seed 42 over C4 (n=1400, seqlen=2048):

• base_c4_n1400_s42_full_2c9a327_calib.pt — clean arm
• ours_c4_n1400_s42_full_2c9a327_calib.pt — traj arm (masked-diffusion activations)
• *_multilingual_n256_s42_* — the multilingual side experiment
• *_manifest.json — per-linear shapes/checksums for each of the above

The two C4 files are a matched pair: same documents, same seed, same order — the only difference is whether activations were collected on clean tokens or on a sampled diffusion trajectory.

Running the missing columns

The three unfinished tasks and their official LLaDA-8B-Base protocol:

Task	few-shot	mode	gen_length / steps / block_length	metric
svamp	5	generation	256 / 256 / 256	exact_match
arc_challenge	0	likelihood (cfg=0.5, mc_num=128, bs 8)	—	acc
math500	4	generation	512 / 512 / 512	exact_match

One cell, from the repo root, with weights/ and code/ in place:

python code/eval/run_lmeval.py \
  --family base --task svamp --arm awq \
  --weights weights/awq/w4g128_c4_traj \
  --num_processes 1 --out_root results/eval_p1/awq_w4g128_c4_traj

python code/analysis/lmeval_to_items.py \
  --family base --task svamp --arm awq \
  --lmeval_dir results/eval_p1/awq_w4g128_c4_traj/awq/lmeval/svamp \
  --out results/eval_p1/awq_w4g128_c4_traj/awq/awq_svamp_items.jsonl

--arm is the family tag (base, ours, asvd, gptq, awq, sparsegpt, wanda, quant_dllm, sink, or ref for dense_bf16), not the arm name; results/eval_p1/.queue/tasks.txt lists the exact <arm_dir>:<task>:<arm>:<weights_path> for all 69 cells. A cell counts as done when results/eval_p1/<arm_dir>/<arm>/<arm>_<task>_items.jsonl is non-empty.

code/jobs/eval_p1_worker.sbatch is a Slurm worker that walks that queue with a directory-based claim protocol, so several GPUs can share the work without coordination. It needs a results/eval_p1/.queue/claims/ directory and paths adjusted to your cluster.

Measured single-GPU cost on an A5000 (an L40S is roughly 2.6× faster): arc_challenge ≈ 3.5 h, svamp ≈ 3.4–6.5 h, math500 ≈ 22 h per arm. lm-eval has no generation cache, so a task killed by a time limit restarts from scratch — give math500 a long allocation.

Environment

transformers with trust_remote_code=True (LLaDA ships custom modelling code), torch, accelerate, and lm-eval. The custom tasks live in code/eval/lm_tasks/ and are passed to lm-eval via --include_path.

Citation & credit

Base model: GSAI-ML/LLaDA-8B-Base. The lm-eval wrapper in code/eval/llada_harness.py derives from the official LLaDA evaluation code (ML-GSAI/LLaDA) and from ML-GSAI/SMDM. Baseline implementations follow the original GPTQ, AWQ, SparseGPT, Wanda, SVD-LLM and ASVD papers; see code/THIRD_PARTY_NOTICES.md where present.