Shankara-A-S/sarvam-30b-distillation-method

Sarvam-30B — Compression & Distillation Method (Notes + Working Pipeline)

A complete, honest engineering write-up of compressing sarvamai/sarvam-30b (SarvamMoE, 19 layers, 128 routed + 1 shared expert, top-6) for the Resilient AI Challenge — Track 1 (lowest inference energy on A100-80GB via stock vllm serve, ≥80% quality recovery per category). This repo documents what worked, what didn't, and every cluster gotcha, plus the reproducible scripts.

TL;DR: W4A16 quantization wins the energy game; brevity distillation does NOT work for this model. Energy = (tokens generated) × (bytes read per token). We maxed the bytes axis (4-bit). The tokens axis (shorten reasoning via distillation) is not reducible here with lightweight LoRA — proven across 3 runs.

---

1. The winning model — W4A16

The submitted model uses a quality-first, serving-safe recipe (merge_quant.py):

• GPTQ (calibrated 4-bit) on attention (query_key_value, dense) + shared experts — the always-active, quality-sensitive paths. Calibration: sarvamai/indivibe (math/STEM/chat/code) + a multilingual builtin fallback, 16 samples × 128 seq (fast; passes all floors).
• RTN (data-free 4-bit) on the 128 routed experts — the sparse bank tolerates it.
• Kept bf16: lm_head (vLLM rejects a quantized lm_head), router gates (routing precision), layer 0.
• Result: ~19 GB, passes every category floor with margin; ~−25% energy vs a standard-W4 cluster.

Why this is the energy optimum: at batch-1 decode the GPU is memory-bandwidth-bound — joules/token ∝ bytes read from HBM. W4 is the practical floor for stock vLLM (sub-4-bit doesn't serve; quantized lm_head fails). 8-bit / FP8 entries read ~2× the bytes ⇒ more energy. GGUF entries don't serve on vLLM at all.

---

2. The distillation attempts (and why they failed)

Goal: cut reasoning tokens (the other energy lever — the model is a "thinking" model that burns hundreds of tokens in <think>…</think>). Three honest attempts, all gated on quality + token count:

Run	Data	SFT	Inference tokens (vs 49,587 shipped)	Floors	Winner
Brevity v1	shortest-of-6, temp 0.9, cap 1536	LoRA attn, gentle	~103k (2.1×)	all pass	❌
Brevity v2	terse, cap 320, temp 0.6 (avg 203 tok)	LoRA attn, LR 2e-5	~99k (2.0×)	all pass	❌
No-think	<think></think> prefill, cap 96 (avg 67 tok)	LoRA attn r=32	~101k (2.0×)	all pass	❌

The finding (consistent across all 3): LoRA-SFT teaches the model to predict terse targets given a terse context (training loss drops cleanly, e.g. 0.34 → 0.26, or no-think 1.7 → 0.65 — it learns), but it cannot override the model's habit of reasoning long when freely generating at temp 1.0. The attention- only adapter (5–10M params) is too weak to change the generation-length distribution; if anything it perturbs the model toward more rambling. Quality stayed perfect every time — only brevity failed.

What would be needed (untested, higher-risk): heavier fine-tuning (higher-rank LoRA on MLP/experts, or full FT) and/or on-policy / RL-style training to actually shift free-generation behavior — which also risks MoE routing stability and the "no finetune after compression" rule (so it must be done before quant).

Legality note (challenge rules, verbatim): distillation is allowed only if the base model is the student; no finetuning after compression. So the legal order is distill → merge → quantize LAST.

---

3. Pipeline (scripts included)

gen_data.py → sft.py → merge_quant.py → bench_offline.py, orchestrated by pipeline.sbatch. guarantee2.sbatch quantizes the base directly (no adapter) and pushes a serving-ready repo.

1. Data-gen (gen_data.py): vLLM offline, best-of-N, keep shortest correct (math/MCQ verified vs known answers; open-ended shortest-coherent). NOTHINK=1 prefills <think></think> to elicit no-think.
2. SFT (sft.py): LoRA on query_key_value,dense only (does NOT touch MoE routing/experts). Custom compute_loss with explicit ignore_index=-100. Stable config: LR 1e-5–2e-5, grad-clip 0.3.
3. Merge + Quant (merge_quant.py): merge adapter on CPU, then GPTQ(attn+shared)+RTN(experts).
4. Bench (bench_offline.py): vLLM offline, energy = ∫ GPU power dt; per-category accuracy + tokens.

---

4. Cluster gotchas (the hard-won list — A100-80GB SLURM, transformers 4.57, vLLM 0.19.1)

These cost the most time; documented so you don't repeat them:

1. download.pytorch.org firewalled → install torch from PyPI (default cu128 works on driver 575).
2. HF repos are Xet-backed → install hf_xet (needs ≥24 GB RAM or it core-dumps); hf_transfer doesn't help Xet repos.
3. trl version conflict → drop it; use transformers.Trainer + peft directly.
4. transformers 4.57 needs tokenizer_class in tokenizer_config (repo lacks it + no AutoTokenizer auto_map) → patch tokenizer_class: PreTrainedTokenizerFast.
5. NCCL multi-GPU "unhandled cuda error" in the cgroup (P2P/SHM/cuMem disable doesn't fix) → vLLM TP=1 for gen/bench. SFT/quant use transformers device_map naive-PP (no NCCL).
6. A ~72 GB per-process GPU memory cap + ~129 GB virtual-address cap (RLIMIT_AS) on the cgroup. A 64 GB model can't load both within them on multi-GPU. Quant: load on CPU, stream layers to GPU.
7. caching_allocator_warmup (transformers 4.57) OOMs on a big contiguous block despite free mem → no-op it for quant (CPU-load); but KEEP it for SFT (device_map GPU load relies on it).
8. RTN's DataFreePipeline.dispatch_model onloads the whole model → OOM on the cap. No-op patch llmcompressor.pipelines.data_free.pipeline.dispatch_model so RTN runs on CPU (it's pure weight math).
9. GPTQ accumulation → pass sequential_targets=["SarvamMoEDecoderLayer"] to oneshot so it onloads/offloads one decoder layer at a time (GPU stays ~5%).
10. A foreign non-SLURM process can pin a GPU SLURM still allocates → pick free GPUs at runtime via torch.cuda.mem_get_info; avoid the busy one.
11. vLLM HTTP server 500s (starlette 1.x + prometheus_fastapi_instrumentator incompat) → bench offline (LLM class), with an if __name__=='__main__' guard (vLLM uses spawn).
12. NFS mmap / RSS limits when loading 64 GB → keep base on home disk (reclaimable page cache), not tmpfs (non-reclaimable, hard-charges the cgroup).

---

4b. Response-economy chat-template trick — DOES NOT transfer to Sarvam (measured)

A config-only "answer concisely / omit nothing / no filler" instruction in the chat template gives non-reasoning models (e.g. Gemma-4-E4B-it) a large token cut (~−44%) with no quality loss. Measured on Sarvam-30B (test_template.py, eval prompts, temp 1.0), the same instruction backfires:

	Total tokens	Math	Questions
Baseline	48,416	1.00	1.00
+ economy instruction	99,823 (2.06×)	0.78 (fails 80% floor)	0.87

Sarvam is a thinking model — the brevity instruction makes it reason more (about how to be economical) inside <think>, doubling tokens, and the meta-reasoning disrupts math. The lever that helps direct-answer models actively hurts reasoning models. Token-axis is closed here via both distillation and template instruction.

5. Honest conclusion

For Sarvam-30B on this challenge, W4A16 (attention+shared GPTQ, experts RTN) is the energy optimum achievable with serving-safe, legal tools. The brevity/token lever — which would compound on top — is not reachable via lightweight LoRA distillation for this model. Most "distilled" competitor entries are 8-bit (higher energy) or quality-recovery rather than token-reduction; the byte axis (4-bit) is where the real energy win lives.

Scripts in this repo are provided as-is for reference. Base model © Sarvam AI.