Skip to content

Support I64 token mappings in safetensors model loading#41

Merged
Pringled merged 2 commits into
MinishLab:mainfrom
maxprogrammer007:main
May 23, 2026
Merged

Support I64 token mappings in safetensors model loading#41
Pringled merged 2 commits into
MinishLab:mainfrom
maxprogrammer007:main

Conversation

@maxprogrammer007
Copy link
Copy Markdown
Contributor

This PR fixes model loading for vocab-quantized models whose mapping tensor is stored as I64 in the safetensors header. The loader previously assumed all mappings were 32-bit integers and decoded the raw bytes with chunks_exact(4), which produced incorrect token ID mappings for models like potion-code-16M.

Changes

  • Decode mapping based on t.dtype() instead of assuming I32
  • Support both I32 and I64 mappings
  • Return a clear error for unsupported mapping dtypes
  • Add a regression test covering both I32 and I64 decoding

Validation

  • cargo test decode_token_mapping --lib
  • cargo test quantized_models_match_float32

fixes #40

@Pringled
Copy link
Copy Markdown
Member

@maxprogrammer007 thanks for fixing this. There was a small formatting issue which caused CI to fail, I've fixed it on your branch. Will release this soon!

@Pringled Pringled merged commit 5d5d425 into MinishLab:main May 23, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

I64 mapping loaded incorrectly

2 participants