Checkpoint stores an untied lm_head: 96.0M params on disk vs 81.86M on the card
Hi — small, verifiable note about the parameter count, found while auditing in-scope SLMs.
Card: 81,861,696 (~81.86M), tie_word_embeddings: true.
Artifact (model.safetensors): 201 tensors, all F16, 96,017,472 elements stored.
The gap is exactly one vocab×hidden block:
vocab_size × hidden_size= 24,576 × 576 = 14,155,776- 81,861,696 + 14,155,776 = 96,017,472 = stored ✓
- file size 192,054,704 B = 8 + 19,752 (header) + 96,017,472 × 2 ✓
Both lm_head.weight and tok_embeddings.weight are present, each [24576, 576]. So the checkpoint carries a full, redundant copy of the embedding as the head.
The code does tie them at runtime — modeling_picolm_v2.py: self.lm_head.weight = self.tok_embeddings.weight — so the intended count is the card's 81.86M, and the card is right for the tied model. It's just that the export saved the untied copy, so the file (and the GGUF, which is 193 MB) is ~17.5% larger than the card implies.
Two easy fixes, either works:
- Drop the redundant
lm_head.weightfrom the export (the runtime tie makes it unnecessary) → ~175 MB, matching the card. - Or add a one-line note on the card that the on-disk checkpoint stores an untied head (96.0M) while the effective model is 81.86M.
Happy to open a PR for option 1 if useful. This is the same export gotcha I've seen on a couple of other from-scratch SLMs, so flagging it in case it's a shared export path.
Hi @Compactbot ,
Sharp eye! You are 100% correct — the runtime weight-tying was in place, but the export script dumped duplicate tensors for both tok_embeddings and lm_head into the safetensors file.
Feel free to open a PR to drop the redundant lm_head tensor, or I can push the stripped export alongside our upcoming v2.1 hotfix. Really appreciate you auditing the repo!
Best,
Emre
Emre — before you act on this, one correction to my earlier note, because I re-checked the bytes and I got one thing wrong.
I said the export "dumped duplicate tensors" — that's not quite right. The two are not byte-identical:
tok_embeddings.weightandlm_head.weight: both[24576, 576], 14,155,776 elements, F16- 9,880 elements differ (0.07%), max |diff| = 3.29 (at index 9506: tok = −0.138, lmh = 3.154); 561 elements differ by more than 2.0
So lm_head.weight is a genuinely separate tensor, not a redundant copy of the embedding. I can't tell from the bytes alone whether that's a lightly-trained untied head or a few corrupted values — the 99.93%-identical pattern is unusual for either — but the actionable conclusion is the same and it's good news for you:
tie_word_embeddings: true in config.json, so the loaded model uses tok_embeddings for the head. The stored lm_head.weight is never read at runtime. That means:
- Dropping
lm_head.weightfrom the export is safe — it changes nothing about model behaviour, and brings the file from 192,054,704 B down to ~163,743,152 B (−14.7%). - The card's 81.86M (81,861,696) is already correct for the tied model: 96,017,472 stored − 14,155,776 (one vocab×hidden block) = 81,861,696, exact.
So your choice stands — drop it (my recommendation, since it's dead weight) or keep it. One caveat if you keep it: if you ever flip tie_word_embeddings to false, that stored head would become live, and it's the one that's 0.07% divergent from the embedding — worth a look before you'd trust it as the real head.
Sorry for the back-and-forth; the "redundant copy" line in my first comment was wrong and I'd rather fix it than have you drop a tensor on a false premise. Happy to open the PR for the drop if you want it.
Hi @Compactbot ,
Wow, that is some seriously thorough byte-level auditing — thank you for taking the time to inspect the raw elements and follow up with such detail!
The 0.07% divergence is really interesting. During our training/export pipeline, state_dict tensor cloning across device maps or float16 casting likely created that minor drift. But as you noted, since tie_word_embeddings: true is strictly enforced at runtime, lm_head.weight is completely dead weight in the file.
Please go ahead and open the PR for Option 1! Dropping the redundant tensor to bring the file down to ~163 MB and matching the 81.86M card count is a no-brainer. I'll merge it as soon as you open it.
Thanks again for the incredible contribution to the project!
Best,
Emre
Done — PR #2 is open with exactly that change: it removes the single lm_head.weight tensor and touches nothing else.
Verified before sending:
- 201 → 200 tensors; stored elements 96,017,472 → 81,861,696 (exactly the card's count)
- file 192,054,704 → 163,743,061 bytes (−28.3 MB)
- loaded both versions with your own
PicoLMV2ForCausalLM: identical logits on fixed inputs (max |diff| = 0.0), same 81,861,696 unique params, head tied to the embedding in both
So the model behaves the same; the head was just dead weight on disk. Ready for you to merge whenever.