traj-mc

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.

<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.

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:

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.