bonsai-qwen3-30b-a3b-teacher-corrected
Qwen3-30B-A3B-Base, full-parameter fine-tuned on a TypeScript retraining pool. This is an intermediate checkpoint from bonsai, a project pruning and distilling this model into a smaller TypeScript-specialized one. It is not pruned and not distilled — it exists to serve as a teacher for later distillation steps, corrected toward the target domain before any student model learns from its logits.
Why this exists
Structured pruning + distillation recipes that distill from the original pretrained model assume access to (something close to) its original pretraining mix. We don't have Qwen3's pretraining mix, so distilling a pruned student directly from the uncorrected base model would hand it teacher guidance drawn from a distribution it was never trained to match downstream. Minitron's practical follow-up (arXiv 2408.11796) calls this "teacher correction": briefly fine-tune the teacher on the target retraining data before using it to distill anything, so the teacher's outputs are still trustworthy on the distribution the student is being trained for.
Training
- Base model:
Qwen/Qwen3-30B-A3B-Base(30.5B total params, MoE, 3B active) - Data: TypeScript source from
bigcode/the-stack-dedup, one epoch, no repeats — 6,000 fresh samples (1500 steps × batch size 4), disjoint from the calibration pool used to score validation loss - Loss: plain LM cross-entropy, no distillation at this stage
- Optimizer: 8-bit AdamW (
bitsandbytes.optim.AdamW8bit) - LR schedule: linear warmup (15% of steps) → cosine decay to a floor; peak LR 2e-5, floor 2e-6 — winner of a 6-config sweep (3 peak LRs × 2 warmup fractions, 150 steps each) on both validation-loss delta and final train loss
- Sequence length: 512 tokens, unpadded (gradient-accumulated single-sample forward/backward, no batching padding waste)
- Hardware: 8x A100 (40GB SXM4),
device_map="auto"with a per-GPU memory cap to leave headroom for gradients/optimizer state, gradient checkpointing enabled - Steps: 1500
Result
| before | after | |
|---|---|---|
| validation loss | 0.9331 | 0.8842 |
Validation loss measured on a held-out calibration shard (disjoint from the retraining pool by construction — shard-level split, not row-level).
What this is not
- Not pruned. Still the full 30.5B-parameter MoE.
- Not distilled from itself — a separate model didn't learn from this one's logits; this model was fine-tuned directly on TypeScript text.
- Not benchmarked on downstream code-generation tasks (HumanEval-TS, MBPP-TS) yet — that comes after distillation into the pruned student candidates this checkpoint will teach.
Reproducing this
Training code: src/bonsai/training/
in the bonsai repo. The exact orchestration script used for this run is
not committed (project convention keeps one-off GPU run scripts as
scratch files, not permanent repo code) — the committed bonsai.training
module is the reusable part; the run itself set peak_lr=2e-5, min_lr=2e-6, warmup_steps=225, total_steps=1500, batch_size=4 and used
bitsandbytes.optim.AdamW8bit as the optimizer.
License
Apache 2.0, inherited from Qwen3-30B-A3B-Base.