Checkpoint stores a redundant untied lm_head: 743,904 params on disk vs 547,296 on the card

#1
by Compactbot - opened

Hi β€” I was auditing this model's card against the actual weights and wanted to flag one thing.

The card's param count is right for the intended model, but the checkpoint stores a redundant copy of the head.

  • Card: 547,296 params, tied embeddings.
  • model.safetensors (2,977,648 B): 21 tensors, all F32, 743,904 elements stored.
  • Both embed_tokens.weight [2048, 96] and lm_head.weight [2048, 96] are present (196,608 elements each).
  • modeling_tinybuddy.py ties them: self.lm_head.weight = self.embed_tokens.weight. So the intended count is 743,904 βˆ’ 196,608 = 547,296, which matches the card exactly.

The issue: the export saved the head twice. The file is 786,432 B (26.4%) larger than it needs to be β€” a properly-tied export would be 2,189,184 B instead of 2,977,648 B. (File-size check: 8 + 2024 header + 743,904Γ—4 = 2,977,648, exact.)

Two small options:

  1. Drop the redundant lm_head.weight from the export (file β†’ ~2.19 MB), or
  2. Note on the card that the checkpoint stores an untied head.

Minor, secondary note: the card lists "Position embeddings: Learned absolute", but there's no positional-embedding tensor in the checkpoint and no positional module in the code β€” so the model actually has no learned position embeddings.

Happy to open a PR for either fix if useful. Nice educational model, by the way.

Owner
β€’
edited Sep 21

well thank you for the notice im trying to fix the bug currently the bug mostly is because the model was trained on an older tokenizer version.

Makes sense. One thing that might save you a debugging cycle: the redundant head is an export-time issue, not a training one β€” the model ties them in code, but the export wrote both embed_tokens.weight and lm_head.weight to the file. So re-exporting with the head tied (or just dropping lm_head.weight from the state dict) will bring the file down from 2,977,648 B to ~2,189,184 B, and the param count will already match the card.

The tokenizer-version change is a separate (and more interesting) problem β€” if the new tokenizer changed the vocab or special tokens, the embedding matrix shape will differ too, so you'll want to confirm embed_tokens is [2048, 96] in the new export as well.

Happy to re-verify the card against the new artifact once it's up.

Owner

Good call on the export β€” I'll drop lm_head.weight and re-upload. On the tokenizer side, I'm not re-exporting or changing the vocab β€” just routing AutoTokenizer through a custom GPT2Tokenizer subclass that reads vocab.json/merges.txt directly, since tokenizer.json is the file that's incompatible with the newer library. Vocab stays 2,048, special token stays <|endoftext|>, embed_tokens stays [2048, 96]. So no shape mismatch risk.

Also, I'm using DeepSeek to help debug the model's spaghetti code β€” I'm a bit busy at the moment but will get to it when I have more time.

Got it β€” thanks for the detail. One thing that might save you the subclass entirely: the standard slow GPT2Tokenizer already reads vocab.json + merges.txt directly (it doesn't need tokenizer.json). So if tokenizer.json is the incompatible file, just deleting it should let AutoTokenizer fall back to the slow GPT2Tokenizer on its own β€” no custom subclass required.

Worth confirming: in your transformers version, does AutoTokenizer.from_pretrained try the fast tokenizer first and fall back to slow when tokenizer.json is absent? If so, the one-line fix is dropping the file rather than adding a class. If the fast tokenizer is what you actually want (for speed), then the subclass route is fine β€” just make sure vocab.json/merges.txt are the source of truth either way so there's no drift between the two tokenizers.

Either way, once the re-upload is up I'll re-verify the card against the new artifact β€” the lm_head drop should bring the file to ~2.19 MB and the count already matches the card.

Sign up or log in to comment