EvanOLeary/laguna-xs2-dense-k8-recon

Laguna-XS.2 → Dense (K=8) · OpenCodeInstruct flavour · reconstruction-pretrained

⚠️ These are PRE-TRAINED weights — Stage 1 only (reconstruction). NOT instruction-tuned.

The dense FFNs have been trained to reconstruct the teacher's MoE outputs (representation alignment). The model can complete code but is not yet a chat/instruct model — that needs the later logit-KD + SFT stages. See Inference behaviour below.

A ~3.0 B fully-dense model densified from poolside/Laguna-XS.2 (33 B / 3 B-active MoE, 256 experts top-8 + shared) by replacing each routed MoE block with a dense SwiGLU FFN, then recovering quality via teacher-forced layer reconstruction on nvidia/OpenCodeInstruct. Method: RADLADS (arXiv:2505.03005) + MoE→Dense (arXiv:2605.28207).

Inference behaviour (real samples from this checkpoint)

✅ Raw code completion works. Prompt def fibonacci(n): →

    if n == 0:
        return 0
    elif n == 1:
        return 1
    else:
        return fibonacci(n-1) + fibonacci(n-2)

(correct recursive Fibonacci — evidence the densification + reconstruction is sound.)

⚠️ Chat/instruct mode is not ready. With the <user>/<assistant> + <think> chat template it degenerates into structural tokens, and some prompts repeat. This is expected for reconstruction-pretraining only — instruct behaviour returns after SFT.

How to use (completion mode)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "EvanOLeary/laguna-xs2-dense-k8-recon"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    repo, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="cuda")

ids = tok("def fibonacci(n):\n", return_tensors="pt").input_ids.to(model.device)
out = model.generate(ids, max_new_tokens=64, do_sample=True, temperature=0.7, top_k=20, pad_token_id=9)
print(tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True))
# (use raw completion, NOT the chat template, until the SFT stage)

Pretraining recipe (detailed & reproducible)

Objective — teacher-forced per-layer feature reconstruction (RADLADS step-1 representation alignment / KRAFTON feature-recon). For each sparse layer ℓ (1..39):

frozen teacher forward  ->  hooks capture   x_ℓ = MoE-block input,  y_ℓ = MoE-block output (target)
student prediction                          pred_ℓ = routed_dense_ℓ(x_ℓ)      # on the TEACHER's input
loss_ℓ = MSE(pred_ℓ, y_ℓ) / mean(y_ℓ²)  +  0.05 * (1 - cos(pred_ℓ, y_ℓ))     # attention-masked
L = mean over all 39 layers (trained in parallel, one backward) ;  update routed_dense only

Teacher-forced inputs mean untrained early layers can't corrupt later layers (no compounding). ÷ mean(y²) (normalize-loss) balances deep vs shallow layers; the cosine term fixes direction.

Stage 0 — DO-ACP warm-start (init, not random): per layer, score the 256 experts by ACP = conditional-prob · √E‖f_e‖², then greedy D-optimal selection — pick the K=8 experts maximizing log det(K_S + λI) of the importance-weighted expert-output Gram K_ij=√(I_iI_j)·G_ij (important AND mutually diverse). Concatenate the 8 experts' gate/up/down into routed_dense (width 8×512=4096), folding α·routed_scaling(2.5) into the down-projection. (−26% deep-MSE vs random.)

Hyperparameters

optimizer	Adafactor, lr 2e-4 (fits all-39-layer on one 80 GB GPU; AdamW state overflows)
seq len / batch	2048 / effective 2 (batch 1 × grad-accum 2)
loss	per-layer normalized MSE + 0.05·(1−cos), attention-masked
trainable	routed_dense only (0.98 B); frozen: attention, embeddings, norms, shared expert, lm_head
steps / tokens	2000 / 8.2 M nominal (1.9 M real; OpenCodeInstruct is short)
dtype / hardware	bf16 / 1× H100 80 GB (~77 GB peak, ~34 min)
data	nvidia/OpenCodeInstruct (100% Python)

Reproduce

# 1) DO-ACP warm-start init
python3 scripts/warm_start_dense.py --output-dir warmstart_student --k 8
# 2) reconstruction pretraining
python3 scripts/train_dense_reconstruction.py \
  --teacher-model poolside/Laguna-XS.2 --student-model warmstart_student \
  --dataset nvidia/OpenCodeInstruct --optimizer adafactor --normalize-loss \
  --seq-len 2048 --batch-size 1 --grad-accum-steps 2 --max-steps 2000 \
  --learning-rate 2e-4 --cosine-weight 0.05 --output-dir out

Results — raw deep-layer MSE (L28-39) vs random-init baseline:

step	random-init	this recipe (warm-start)
100	0.073	0.036
1000	0.032	0.023
2000	0.032	0.022

Training curves (OpenCodeInstruct, 2000 steps)

Left: normalized total loss 0.7→~0.25. Middle: per-depth reconstruction MSE — shallow L1-10 solve to ~1e-4, deep L28-39 sit at ~2e-2 (the hard layers). Right: warm-start vs random-init deep-MSE (−26%).

Per-layer MSE (log10): deep layers (bottom, layer ~30 hottest) start high and cool over training; shallow layers (top) are easy.

Architecture

~3.0 B dense (laguna_dense): 1 SwiGLU FFN (width K8×512=4096) + kept shared expert per layer; attention (48/8 GQA, 30 SWA + 10 global), embeddings, norms copied from the teacher. Hidden 2048 · 40 layers · 262 k ctx · 100 352 vocab · SiLU. 5.99 GB bf16.

Roadmap (what makes it usable)

1. (done) Stage-1 reconstruction pretraining ← this checkpoint.
2. Kernel-mixture variant (KernelBook + CUDA + OpenCode).
3. Logit-KD (KL to teacher) — recover function-level fidelity.
4. SFT on code/instruct — recover chat + robustness (fixes the degeneration above).

Limitations

Pre-trained only — not instruction-tuned, not robust, chat template unusable yet. Use for code completion / research, not production. Reconstruction MSE is a proxy; perplexity + KernelBench are the end-to-end metrics.

Refs: RADLADS arXiv:2505.03005 · MoE→Dense arXiv:2605.28207. Poolside Laguna XS.2 hackathon.