Forward revision to the config, weight and tokenizer loads - #4222
Conversation
- Fix #3544: Add revision parameter to AutoConfig, AutoModelForCausalLM, AutoModelForSequenceClassification, and load_correct_tokenizer calls in FastLlamaModel.from_pretrained. This enables loading specific model revisions/branches from HuggingFace Hub. - Fix #3667: Escape single quotes in system messages before substituting into Jinja2 templates. This prevents TemplateSyntaxError when system messages contain apostrophes (e.g., "user's" in Vicuna templates). Signed-off-by: majiayu000 <1835304752@qq.com> (cherry picked from commit b0a6e41)
- Add revision to load_vllm_kwargs in llama.py to fix config/weights mismatch - Add revision to PEFT AutoConfig calls in loader.py (FastLanguageModel & FastModel) Addresses reviewer feedback from @chatgpt-codex-connector and @Datta0 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> (cherry picked from commit 14f89e4)
Propagate revision parameter to all from_pretrained calls in vision.py to ensure consistent version pinning for vision models. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com> (cherry picked from commit c5aa4ec)
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request resolves two critical issues: enabling the use of specific model revisions during loading and preventing Jinja2 template parsing errors caused by unescaped single quotes in chat templates. These changes enhance the robustness and flexibility of model and tokenizer loading, and ensure proper rendering of chat templates. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
The pull request enhances model and tokenizer loading by introducing a revision parameter across various from_pretrained calls in unsloth/models/llama.py, unsloth/models/loader.py, unsloth/models/vision.py, and unsloth/tokenizer_utils.py, allowing users to specify a particular version (e.g., branch, tag, or commit hash) when loading resources from Hugging Face. Additionally, unsloth/chat_templates.py was updated to escape single quotes in system messages, preventing potential Jinja2 template syntax errors.
This comment was marked as low quality.
This comment was marked as low quality.
The branch has drifted nine months. Resolving the merge to main's content
everywhere, because every hunk it carried is now either landed, superseded
or actively harmful:
- chat_templates.py escaping: superseded. main no longer uses re.sub, and
_escape_jinja_literal already covers backslash, both quotes and CR.
- loader.py AutoConfig(revision = ...): already on main.
- vision.py adding revision to the FastBaseModel.from_pretrained signature:
harmful. revision arrives via **kwargs there, and binding it as a named
parameter drops it from the weight load and from kwargs.get("revision").
- llama.py load_vllm(revision = ...): raises TypeError. load_vllm has no
revision parameter and load_vllm_kwargs is not filtered.
The revision fix itself follows in the next commit.
FastLlamaModel.from_pretrained took a `revision` argument and never read it, so the config, the weights and the tokenizer all came from the repo's default branch while the caller believed they had pinned a ref. Reported in #3544 by someone versioning their fine-tunes with branches, which makes it a silently wrong base checkpoint rather than an error. Forward it in llama.py (both AutoConfig loads, the three model loads, the tokenizer, the prefetch warm and the fp8 scale restore), plumb it through load_correct_tokenizer, and read it from kwargs in vision.py for the four AutoConfig, two processor and two tokenizer loads plus the VLM processor fallback. vision.py must not bind it as a named parameter: the weight load there forwards **kwargs, so binding it would drop it from that load. model_name is not always the repo the caller named. get_model_name can swap in a pre-quantized mirror, _offline_quantize_to_fp8 an fp8 temp dir, ModelScope a local snapshot, and fast_inference_setup a -bnb-4bit variant, and use_exact_model_name only gates the first of those. A ref from the original repo does not exist on the substitute, so _revision_for_resolved_repo drops it with a warning naming both repos when the resolution changed the name. The adapter load keeps the caller's revision, since that one really is for old_model_name. Supersedes the earlier attempt on this branch, whose chat-template hunk is handled by #7731 and #7746, whose vision.py signature change caused the drop described above, and whose load_vllm(revision = ...) raised TypeError because load_vllm has no such parameter. Fixes #3544
for more information, see https://pre-commit.ci
Four fixes from review: - The gate ran after the AutoConfig and PeftConfig probes, which already used the raw revision against the resolved name, so a pinned load_in_4bit load failed against the mirror instead of warning. Gate right after the resolution block and point both probes at the gated value, then re-gate before dispatch for the later fast_inference_setup remap. Feeding the second call the first result keeps the warning to one. - On a PEFT load model_name is necessarily the base model, so the late gate warned "Ignoring revision" for every versioned adapter and told the caller to pass use_exact_model_name, which cannot stop an adapter resolving its base. Skip the late gate for PEFT; PeftModel.from_pretrained already loads the adapter with the caller's revision. - load_vllm takes no revision, so vLLM fetches the default branch. Pinning only the config and the tokenizer put two refs in one model, which is worse than the old behaviour of ignoring the revision outright. Drop the pin with a warning before the config load whenever vLLM owns the weights. - _hub_repo_or_local_path resolved a cached snapshot without the revision, so an offline or local_files_only tokenizer load silently got the default ref: a revision handed to from_pretrained cannot re-point a local directory. Thread it into _resolve_hub_repo_local_dir and both call sites. Five new tests, one per fix, all failing before it.
for more information, see https://pre-commit.ci
The vLLM guard sat at the end of the same block that turns fast_inference off when vLLM is missing or the GPU is older than sm70. In that case the load falls through in-process and can honour the revision, but the guard dropped it anyway. Re-check fast_inference in the condition.
Three more from review: - A num_labels load goes through AutoModelForSequenceClassification in-process no matter what fast_inference says, so the vLLM guard was discarding a revision the load could have used. Condition it on the same `fast_inference and num_labels is None` predicate the prefetch warm already uses. - use_exact_model_name only gates the mapper substitution. The ModelScope download, the ALLOW_PREQUANTIZED_MODELS strip and fast_inference_setup ignore it, so the warning was sending callers round the same loop. Record whether the mapper is what moved the name and only offer the remedy then. - The tokenizer does not always come from the base model's repo. Loading a PEFT repo with an explicit tokenizer_name pointing at the adapter dropped the pin for the tokenizer while PeftModel loaded the adapter from the requested ref, mixing two refs. _revision_for_tokenizer_repo now resolves it where the repos are known and both dispatches carry it, replacing the tokenizer_name == model_name guess in llama.py and vision.py. vision.py pops it from kwargs, since the weight load forwards **kwargs and transformers has no such argument. Seven new tests, all failing before this.
for more information, see https://pre-commit.ci
Three more from review, all fallout from splitting tokenizer_revision out: - Skipping the late gate for PEFT leaves base_revision naming the adapter, and a remote PEFT load without an explicit tokenizer_name reads its tokenizer from the base repo, so that ref was handed to the wrong repository. Both dispatches now derive one model_revision and pass it to the base load and to the tokenizer resolution alike, so the base tokenizer can only ever get the base model's ref. - FastLlamaModel is exported, and the architecture wrappers forward `revision` through **kwargs without the new internal tokenizer_revision, so a direct call pinned the config and weights while the tokenizer read the default branch. Fall back to `revision` when the tokenizer repo is the model repo, before the warm so it does not fetch the wrong ref either. - The vLLM guard cleared only the model pin, leaving vLLM on the default branch with the tokenizer still on the requested ref. Clear both, in llama.py and in the parallel FastBaseModel block. Seven new tests.
for more information, see https://pre-commit.ci
FastModel withheld the probed config from FastBaseModel on the vLLM path, but model_types, auto_model and the text-only decision had already been derived from it, so default-branch weights could load with pinned-ref dispatch. The drop now happens before the probe instead, using the same predicate FastBaseModel does, which makes that guard a no-op on this path and lets the config go down untouched again. FastLanguageModel keeps its drop inside llama.py: that one also turns fast_inference off on pre-Volta GPUs and for a num_labels load, and the loader cannot see either without duplicating the device checks, so gating early there would discard a pin llama.py would have honoured. The fp8 cache name sanitized the ref by replacing every unsafe character with the same one, so release/v1 and release.v1 shared a directory and the second load reused the first ref's artifact. A digest of the raw ref now rides along with the readable form.
for more information, see https://pre-commit.ci
…he saved ref FastLanguageModel still probed the config at the pinned ref while llama.py dropped that same ref for its vLLM load, so model_types could pick the architecture class off one ref and load weights from another. It now drops the pin before the probe like FastModel does, through _vllm_will_load_weights in llama.py, which llama.py itself now calls: the language path also falls back in-process on pre-Volta GPUs and for a num_labels load, so the predicate has to live where those checks are rather than be guessed at by the loader. That drop runs before is_peft is known, and it was zeroing the ref the PeftConfig probe reads. An adapter is loaded in-process by peft, so it keeps the ref: adapter_revision holds the value from before the vLLM drop. Pinning the tokenizer also desynced the save path, which restores tokenizer.model from tokenizer.name_or_path and so had no idea which branch to read. The loaded ref is now stamped on the tokenizer the way local_files_only and cache_dir already are, and the sentencepiece probe, its memo key and the restore all use it.
for more information, see https://pre-commit.ci
FastBaseModel builds its processor without going through load_correct_tokenizer, so the stamp save.py reads was only being applied on the text path and a pinned FastVisionModel load still restored tokenizer.model from the default branch. Stamped at the return rather than at each of the processor branches, so the AutoTokenizer fallback that runs when patch_tokenizer raises cannot lose it either.
|
Codex Review: Didn't find any major issues. Another round soon, please! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: eff8c12973
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if base_revision is not None and _vllm_will_load_weights( | ||
| fast_inference, kwargs.get("num_labels") | ||
| ): | ||
| logger.warning_once( | ||
| f"Unsloth: Ignoring revision = `{base_revision}` since vLLM loads weights " | ||
| "from the default branch. Use `fast_inference = False` to load a pinned revision." | ||
| ) | ||
| base_revision = None |
There was a problem hiding this comment.
Preserve the adapter ref during repository-type detection
When fast_inference=True loads a PEFT adapter from a non-default revision, clearing base_revision here makes the subsequent AutoConfig probe inspect the adapter repository's default branch while PeftConfig inspects the requested revision. If those refs have different layouts—for example, the requested ref is adapter-only while the default branch contains an incompatible model config—the unpinned AutoConfig probe can raise early or, on older supported Transformers versions, make the repo appear to contain both a model and an adapter. vLLM only owns the eventual base-model weight load, not this adapter-repository classification, so both probes should use the adapter ref until is_peft is known; the same ordering is duplicated in FastModel.from_pretrained.
Useful? React with 👍 / 👎.
Fixes #3544. Picks up the branch @majiayu000 opened as a replacement for #3793 and brings it to current
main; their commits are kept in the history.FastLlamaModel.from_pretraineddeclared arevisionparameter and never read it. The config, the weights and the tokenizer all loaded from the repo's default branch while the caller believed they had pinned a ref, so the failure mode is a silently wrong base checkpoint rather than an error. From the issue:unsloth/models/loader.pyalready forwardedrevisiondown to the dispatch, so this was the last mile.What now honours
revisionmodels/llama.pyAutoConfigloads, the causal-LM and sequence-classification loads,load_correct_tokenizer, the snapshot prefetch, the fp8 scale restoremodels/vision.pyAutoConfigloads, both processor loads, both tokenizer loads, the VLM processor fallbacktokenizer_utils.pyload_correct_tokenizer/_load_correct_tokenizergain the parameter and pass it to the slow and fast tokenizer loadsmodels/loader_utils.py_load_pretrained_tokenizer_fastTwo places deliberately do not take it, and both have a comment saying why:
auto_model.from_pretrainedinvision.pyalready receives it through**kwargs, so an explicit keyword there is a duplicate-keywordTypeError. For the same reasonFastBaseModel.from_pretrainedmust not bindrevisionin its signature: that would remove it from**kwargsand drop it from the weight load. There is a test pinning this.The part that is not mechanical
model_nameis not always the repo the caller named. Before dispatch it can become a pre-quantized mirror (get_model_name), a local fp8 temp dir (_offline_quantize_to_fp8), a ModelScope snapshot, a-bnb-4bitstrip, the PEFT base model, or a-unsloth-bnb-4bitto-bnb-4bitswap (fast_inference_setup).use_exact_model_name = Trueonly gates the first of those.A ref from the original repo does not exist on the substitute, so pinning it there would 404, or worse resolve a same-named branch on a different repo.
_revision_for_resolved_repotherefore drops the revision when resolution changed the name, and says so:The issue's own case is unaffected by this: a user's own repo is never in the mapper tables, so nothing is remapped and the revision is honoured. No load that works today changes behaviour, it just stops being silent when it ignores you. The adapter load keeps the caller's
revision, since that one really does belong toold_model_name.What was dropped from the original branch
The first commit here resolves the nine-month drift by taking
main's content, because each of the earlier hunks had gone stale:chat_templates.pyquote escapingmainno longer usesre.subthere, and #7731 / #7746 landed_escape_jinja_literal, covering backslashes, both quote characters and CR.loader.pyAutoConfig(revision = ...)main.vision.pyaddingrevisionto the signaturellama.pyload_vllm(revision = ...)load_vllmhas norevisionparameter andload_vllm_kwargsis not filtered, so this raisedTypeErroron every vLLM load. There is now a test pinning that too.Tests
tests/python/test_revision_forwarding.py, 20 cases, AST-structural like the existingtest_prefetch_snapshot_scope.pyandtest_fast_language_model_text_only.py, so no GPU, no network and no gated checkpoint. It asserts every load listed above carriesrevision, thatFastBaseModeldoes not bind it, thatload_vllm_kwargsdoes not contain it, and execs_revision_for_resolved_repoout of the AST to cover pass-through, the remap drop over three real remap shapes, and the warning naming both repos.Against
main18 of the 20 fail; the 2 that pass are the regression guards. Locally:Repo tests (CPU)3806 passed 0 failed, ruff clean, andscripts/verify_import_hoist.pyreports no blockers on all five changed files.