yuchengm/quadtok-gen-700m-2level

QuadTok generator — 775.2M, 2-level, 400k steps (gen_700m_local)

Causal autoregressive QuadtreeGPT over QuadTok content-adaptive quadtree tokens, class-conditional on ImageNet-1k at 256px. Trained to completion: 400,000 steps, exit 0, zero errors or OOMs.

READ THIS BEFORE LOADING — freqs_cis will hard-error otherwise

These weights were produced by a fork whose QuadtreeGPT uses seq_len=1408 (3-level capacity) rather than upstream's 512. freqs_cis is a persistent RoPE buffer of length 1 + 2*seq_len, so this checkpoint carries a (2817, 32, 2) table where upstream code builds (1025, 32, 2). load_state_dict treats a size mismatch as an error even under strict=False:

RuntimeError: Error(s) in loading state_dict for QuadtreeGPT:
  size mismatch for freqs_cis: copying a param with shape torch.Size([2817, 32, 2]) from
  checkpoint, the shape in current model is torch.Size([1025, 32, 2]).

freqs_cis is derived, never learned — a pure function of head_dim and length — and it is a strict prefix (verified: torch.equal(ckpt["freqs_cis"][:1025], model.freqs_cis) is True). So drop it and let __init__'s own table stand:

meta = torch.load("ckpt/step_400000/meta.pt", map_location="cpu")
sd = {k: v for k, v in meta["ema"].items() if k != "freqs_cis"}
msg = model.load_state_dict(sd, strict=False)
assert not msg.unexpected_keys and set(msg.missing_keys) <= {"freqs_cis"}

Verified end to end: step=400000 dropped_derived=['freqs_cis'] missing=1 unexpected=0, 775.2M params. RoPE is unaffected — positions are indexed, so the extra rows are never read.

Also: mlp_ratio=1. SwiGLU FeedForward already applies the 2/3·4·dim rule, so passing 4 double-counts it. Load with mlp_ratio=1 (the older Google-Drive reference checkpoint needs 4; these weights do not).

Contents

ckpt/step_{25000..400000}/, 20 checkpoints, 220 files, 231.1 GB:

file	size	what
meta.pt	2.9 GB	{"step": int, "ema": state_dict} — EMA weights, use these for inference
state/model.safetensors	2.9 GB	raw (non-EMA) weights
state/optimizer.bin	5.8 GB	AdamW moments — only needed to resume
state/random_states_0..7.pkl	16 KB each	dataloader/RNG state per rank

Kept: every 25k to 375k, then 390k/392.5k/395k/397.5k/400k.

Recipe

model	xlarge — embed 1280 / depth 36 / heads 20, 775.2M (LlamaGen-XL dims)
vocab	16384, token_size 8
tree	2-level, num_patch_side_list [1,2,4,8,16] (LOD3 8×8 + LOD4 16×16), ≈231 tok/img
data	yuchengm/quadtok_data :: pretok_2level, 1470 shards / 24 GB
batch	232,448 tok/step = 116,224 tok/GPU-group × accum 2 ≈ 1000 img/step
optim	AdamW β(0.9, 0.95), wd 0.05, clip 0.5
lr	4e-4 → 2e-5 cosine, warmup 50,000
steps	400,000
EMA	0.9999
precision	bf16, no grad checkpointing, flash_attn_varlen_func (block-diagonal causal)
hardware	8× H200, 44.0 GB peak/GPU
wall clock	3 d 20 h 26 m (up to 302k tok/s, 0.76 s/step once the node was uncontended)

Training curve

Cross-entropy in nats over the 16384-way codebook (ln(16384) = 9.7041 is the uniform floor).

step	loss	token acc
0	9.7041	0.0000
25000	7.3182	0.0289
50000	7.0928	0.0353
75000	6.9412	0.0408
100000	6.8411	0.0444
125000	6.7801	0.0467
150000	6.7495	0.0481
175000	6.6984	0.0501
200000	6.6225	0.0538
225000	6.6021	0.0544
250000	6.5351	0.0575
275000	6.4604	0.0616
300000	6.4109	0.0639
325000	6.3240	0.0680
350000	6.2958	0.0700
375000	6.2558	0.0727

Final: loss 6.2358, token accuracy 7.29% → perplexity ≈ 513 of 16384.

For reference, the measured entropy floors on this data: H(code) = 9.6969 and H(code | lod) ≈ 9.67 — the codebook is near-uniform (all 16384 codes used, most frequent 0.033%), so the unigram "free lunch" is only ~0.03 nats. The ~3.4 nat gap below the conditional unigram floor is genuine conditional modeling.

Token order — the trap that metrics cannot catch

AR training requires tokens in _get_ordered_nodes order (BFS grouped by LOD), not slot/spatial order. Reconstruction metrics are order-insensitive and will not detect a wrong order; an earlier run with slot order reached CE 7.63 yet FID 45. The structural test is that lod_indices must be non-decreasing within every sample. The pretok_2level data used here passes it: 0 non-monotone samples across all 1470 shards, mean L = 231.1, max 316, lod ∈ {3, 4}, max_patch_idx = 255.

Caveat

FID has not been computed for these weights yet — the loss curve is the only evidence of quality here. Decoding to images additionally requires the matching QuadTok tokenizer (VQ decoder), which is not in this repo.