Skip to content

fix: upload tokenizer/config files with hub checkpoints - #656

Open
salch-cred wants to merge 3 commits into
PsycheFoundation:mainfrom
salch-cred:fix/checkpoint-extra-files
Open

fix: upload tokenizer/config files with hub checkpoints#656
salch-cred wants to merge 3 commits into
PsycheFoundation:mainfrom
salch-cred:fix/checkpoint-extra-files

Conversation

@salch-cred

Copy link
Copy Markdown

Summary

When a run's init checkpoint comes from one HF repo but checkpoints are uploaded to a different repo, the checkpoint repo only ever received the safetensors files — it was missing added_tokens.json, config.json (no— that one was sent), merges.txt, special_tokens_map.json, tokenizer.json, tokenizer_config.json, vocab.json and friends, so it was not loadable as a standalone model repository.

Fixes #631.

There were two gaps, fixed in two commits:

1. The download path never fetched these files.
MODEL_EXTENSIONS in both hub.rs and gcs.rs was [".safetensors", ".json", ".py"], so merges.txt, tokenizer.model (sentencepiece), and chat_template.jinja never made it to local disk in the first place. The two duplicated constants are consolidated into one MODEL_FILE_EXTENSIONS constant in file_extensions.rs (alongside the existing DATA_FILE_EXTENSIONS), extended with .txt, .model, .jinja.

2. The checkpoint extra-files filter missed several file names.
The checkpoint_extra_files filter in init.rs (duplicated in the hub and GCS branches) only carried config.json, tokenizer.json, tokenizer_config.json, special_tokens_map.json, generation_config.json, and .py files. Both duplicates are extracted into one is_checkpoint_extra_file helper carrying the full list from the issue (added_tokens.json, merges.txt, vocab.json) plus tokenizer.model and chat_template.jinja.

Matching is now by exact file name (via file_name()) instead of the previous path-suffix ends_with, which incidentally avoids false positives like a file named mytokenizer.json matching the old tokenizer.json suffix test.

Behavior change

Model downloads now also fetch tokenizer artifacts (a few small text/binary files — merges.txt is typically ~456KB, tokenizer.model ~500KB for a 32k vocab) and re-upload them with every checkpoint. Every caller of download_model_repo_async/_sync and download_model_from_gcs_* benefits with no API change.

Notes for reviewers

  • All call sites of the old constants are updated (hub.rs x2, gcs.rs x1); MODEL_FILE_EXTENSIONS is also re-exported from the crate root next to the existing re-exports.
  • The remaining .py match is kept for custom modeling code repos.
  • Local checkpoint dirs (model_is_local branch reading an unpacked local repo) flow through the same helper, so a local init repo now also contributes its tokenizer files.

The hub and GCS download paths only fetched .safetensors, .json and .py
files, so tokenizer artifacts like merges.txt, vocab.json, added_tokens
and tokenizer.model (sentencepiece) were never present locally and
could not be re-uploaded with checkpoints. Consolidates the duplicated
MODEL_EXTENSIONS lists into a shared MODEL_FILE_EXTENSIONS constant and
extends it with .txt, .model and .jinja.

Part of PsycheFoundation#631
The checkpoint extra-files filter missed added_tokens.json, merges.txt,
vocab.json and tokenizer.model, so checkpoint repos for runs whose init
checkpoint differs from the checkpoint repo ended up with weights only
and were not loadable as standalone model repos. Extracts the duplicated
filter in both the hub and GCS init branches into one helper with the
full file list (also adds chat_template.jinja and tokenizer.model).

Fixes PsycheFoundation#631
@salch-cred

Copy link
Copy Markdown
Author

Note: I just noticed #647 addresses the same issue (#631) with a similar approach — apologies for the overlap, I searched open PRs for \checkpoint/\ okenizer\ but it didn't surface.

Differences vs #647, for the maintainers' convenience:

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.

a checkpoint hf repo != the init repo does not upload all the original repo fiules, *only* the safetensors

1 participant