fix: upload tokenizer/config files with hub checkpoints - #656
Open
salch-cred wants to merge 3 commits into
Open
Conversation
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
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:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
safetensorsfiles — it was missingadded_tokens.json,config.json(no— that one was sent),merges.txt,special_tokens_map.json,tokenizer.json,tokenizer_config.json,vocab.jsonand 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_EXTENSIONSin bothhub.rsandgcs.rswas[".safetensors", ".json", ".py"], somerges.txt,tokenizer.model(sentencepiece), andchat_template.jinjanever made it to local disk in the first place. The two duplicated constants are consolidated into oneMODEL_FILE_EXTENSIONSconstant infile_extensions.rs(alongside the existingDATA_FILE_EXTENSIONS), extended with.txt,.model,.jinja.2. The checkpoint extra-files filter missed several file names.
The
checkpoint_extra_filesfilter ininit.rs(duplicated in the hub and GCS branches) only carriedconfig.json,tokenizer.json,tokenizer_config.json,special_tokens_map.json,generation_config.json, and.pyfiles. Both duplicates are extracted into oneis_checkpoint_extra_filehelper carrying the full list from the issue (added_tokens.json,merges.txt,vocab.json) plustokenizer.modelandchat_template.jinja.Matching is now by exact file name (via
file_name()) instead of the previous path-suffixends_with, which incidentally avoids false positives like a file namedmytokenizer.jsonmatching the oldtokenizer.jsonsuffix test.Behavior change
Model downloads now also fetch tokenizer artifacts (a few small text/binary files —
merges.txtis typically ~456KB,tokenizer.model~500KB for a 32k vocab) and re-upload them with every checkpoint. Every caller ofdownload_model_repo_async/_syncanddownload_model_from_gcs_*benefits with no API change.Notes for reviewers
hub.rsx2,gcs.rsx1);MODEL_FILE_EXTENSIONSis also re-exported from the crate root next to the existing re-exports..pymatch is kept for custom modeling code repos.model_is_localbranch reading an unpacked local repo) flow through the same helper, so a local init repo now also contributes its tokenizer files.