[kashif]

What does this PR do?

When a trainer loads a model from a pinned revision through model_init_kwargs, its automatically created processing class now uses the same revision. This prevents model weights from being paired with a tokenizer, processor, or chat template from the repository default branch.

The fix covers the seven CLI-backed main trainers (SFT, DPO, KTO, GRPO, RLOO, Reward, and Distillation) and every duplicated experimental path (A2PO, SDFT, SDPO, TPO, CPO, BCO, ORPO, SSD, GOLD, IW-OPD, Async GRPO, and Async Distillation). It also forwards the effective revision to string-loaded reward-model tokenizers in GRPO, RLOO, SDPO, and Online DPO, and to GOLD teacher processors/tokenizers.

Only revision is forwarded from model initialization kwargs. Other values such as token, cache_dir, and subfolder remain out of scope. When a model is already instantiated, the processing-class revision remains unset because model initialization kwargs are ignored on that path.

The regression test uses the existing pinned revision 8913f5819566 of trl-internal-testing/tiny-Qwen2ForCausalLM-2.5 and checks the resulting chat template, reproducing the actual bug without mocks.

Fixes #6974

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that is the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that is the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

AI writing disclosure

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

Anyone in the community is free to review the PR once the tests have passed.

---

Note

Medium Risk
Wide but mechanical init-path change across many trainers; mis-forwarding revision could break Hub loads or tokenization, though behavior is gated on string model loading and GOLD’s cross-repo guard.

Overview
Trainers that auto-load a tokenizer or processor from a Hub model id now pass the same revision used in model_init_kwargs (or GOLD’s teacher init kwargs), so pinned weights are not paired with default-branch chat templates or vocab.

The change is applied consistently across main trainers (SFT, DPO, KTO, GRPO, RLOO, Reward, Distillation) and experimental trainers, with config help text updated to document the behavior. Revision is only forwarded when the model is loaded from a string; pre-instantiated models still leave revision unset for the processing class.

GOLD additionally pins teacher processors/tokenizers to the teacher revision, but drops that revision for ULD when teacher_tokenizer_name_or_path points at a different repo (avoiding invalid cross-repo commit refs). GRPO, RLOO, SDPO, and Online DPO apply the same revision when auto-loading reward-model tokenizers from string reward checkpoints.

Tests assert the pinned Qwen2 revision’s chat-template length on SFT init and cover GOLD ULD same-repo vs cross-repo teacher tokenizer behavior.

^{Reviewed by Cursor Bugbot for commit 62560d9. Bugbot is set up for automated code reviews on this repo. Configure here.}

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-08-29T21:52:25.134102Z	332c18b	PR opened

ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 332c18b3eb

ℹ️ 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".

[albertvillanova]

Thanks for picking this up. The mechanism is right and I checked all 11 touched trainers on a6f3a181a: every path assigns model_revision before use, and the semantics are correct in the non-obvious case. Forwarding the revision only when the trainer itself loaded the model from a string is the right choice, since model_init_kwargs is documented as ignored for an already-instantiated model. The tempting one-liner (args.model_init_kwargs or {}).get("revision") (the idiom currently sitting in gold_trainer.py:806) would silently apply a stale revision there. Also revision=None passed explicitly is a no-op, cached_file and hf_hub_download both resolve None to main, so there is no behavior change for existing users.

One point in favour of this PR that is worth adding to the description: the seven main trainers you fixed are exactly the seven CLI entry points (trl/scripts/{sft,dpo,kto,grpo,rloo,reward,distillation}.py) that each build model_init_kwargs = dict(revision=model_args.model_revision, ...). So today trl sft --model_revision <sha> loads weights from <sha> and the tokenizer from main. This fixes a documented CLI flag, not just a niche API path.

A few things before this lands.

1. Eight copies of the same pattern are still unfixed

The Codex review named tpo, sdft, sdpo and a2po, and those are now handled, but that list was incomplete. On current main the identical pattern (string model loaded with model_init_kwargs, then processing_class auto-loaded from the same repo id without revision) is still live in:

File	Line
trl/experimental/cpo/cpo_trainer.py	290
trl/experimental/bco/bco_trainer.py	606
trl/experimental/orpo/orpo_trainer.py	300
trl/experimental/ssd/ssd_trainer.py	160
trl/experimental/gold/gold_trainer.py	810
trl/experimental/iw_opd/iw_opd_trainer.py	438
trl/experimental/async_grpo/async_grpo_trainer.py	830
trl/experimental/async_distillation/async_distillation_trainer.py	963

That is 4 of 12 experimental copies fixed. Two are self-evidently wrong in place:

• iw_opd_trainer.py:533 already does teacher_revision = teacher_model_init_kwargs.get("revision", args.teacher_model_revision) and forwards it to the teacher tokenizer at :537. The student tokenizer 100 lines above does not. Same file, opposite behavior.
• gold_trainer.py:806 computes self.model_revision = (args.model_init_kwargs or {}).get("revision") and then never uses it anywhere in the repo (single occurrence). The processing-class load that needs it is four lines below at :810.

Not affected, so correctly untouched: online_dpo (raises when processing_class is None), minillm and prm (no auto-load), gspo/papo (subclass GRPO and inherit the fix).

Either finish the sweep or state in the description that experimental coverage is partial and open a follow-up. As it stands the PR reads as complete when it is not, and per the consistency rule in AGENTS.md a partially propagated fix is the failure mode we most want to avoid.

2. model_init_kwargs help text is now inaccurate

TRL is already explicit about this for the neighbouring field: the trust_remote_code help in grpo_config.py:436-439 says "Forwarded to AutoModelForCausalLM.from_pretrained and AutoProcessor.from_pretrained". The model_init_kwargs help directly above it still mentions only AutoModelForCausalLM.from_pretrained, which is wrong for revision after this change. Worth updating in the affected configs (docstring and metadata["help"]).

3. Only revision is forwarded, which is fine but should be said

model_init_kwargs can equally carry token, cache_dir, local_files_only and subfolder, all of which affect where the processor is resolved from just as much as revision does. A private repo loaded with model_init_kwargs={"token": ...} still fails on the auto-processor load. The issue asked for the revision "at minimum", so revision-only is a defensible scope, but it is worth saying so explicitly, and token looks like a natural follow-up.

4. The test helper is heavier than it needs to be

assert_processing_class_revision plus wrap_classmethod_from_pretrained adds 67 lines to shared tests/testing_utils.py, patches four classes (AutoProcessor, AutoTokenizer, PreTrainedTokenizerBase, ProcessorMixin) through a hand-rolled classmethod re-wrapper, and carries an 11-line docstring about classmethod descriptor binding.

It is correct, I traced it: mock.patch.object reads cls.__dict__ first so the raw classmethod is restored, and the wrapper calls original_func(cls, ...) with the real call-time cls. But the complexity is self-inflicted: it only arises because the two base classes are patched as well. Each trainer calls exactly one of AutoProcessor / AutoTokenizer, and the trainer source tells you which, so a short inline spy does the same job:

with patch.object(AutoProcessor, "from_pretrained", wraps=AutoProcessor.from_pretrained) as spy:
    SFTTrainer(model=..., args=SFTConfig(..., model_init_kwargs={"revision": "main"}), train_dataset=dataset)
assert spy.call_args.kwargs["revision"] == "main"

Given the "leanness and simplicity above all" guidance in AGENTS.md, I would drop the helper. If it stays, wrap_classmethod_from_pretrained should be private, it has exactly one caller.

Smaller test notes:

• expected_revision="main" is the default revision, so the assertion covers kwarg plumbing rather than which files are actually fetched. It does fail before the fix (None != "main"), so it is a valid regression test. A stronger one would pin a non-default revision on a trl-internal-testing fixture whose chat template differs and assert the resulting template, which is the failure the issue describes. Optional, and it needs a Hub fixture first.
• report_to="none" is missing in 8 of the 11 new tests (present only in a2po, sdft, sdpo). Harmless at init time, but it goes against the convention in these files.
• If the per-trainer tests stay, they should cover the remaining experimental trainers too once point 1 is addressed, otherwise the suite advertises coverage it does not have.

Out of scope, for a separate issue

trl/generation/vllm_generation.py:348 calls LLM(model=model.name_or_path, ...) with no revision, so in colocate mode vLLM resolves config and tokenizer from the default branch. Weights are synced from the training model, so the exposure is a config/tokenizer mismatch rather than wrong weights. Not for this PR, but same bug class and worth tracking.

Happy to approve once the propagation in point 1 and the help text in point 2 are sorted. The Codex thread above is stale, it reviewed 332c18b and a6f3a181a addressed it, so a reply there would clear it.

[qgallouedec]

Trainer side is right.

I'd drop the tests though. assert_processing_class_revision is 67 lines patching four classes with a hand-rolled classmethod re-wrapper, and what it buys is an assertion that revision="main" comes back out after you passed revision="main" in. That only fails if someone deletes the kwarg, and nothing about the loaded processor differs, so it can't catch the thing the bug actually caused, which is the wrong chat template.

Testing it properly needs a fixture repo with two revisions whose tokenizers differ, then it's a four line test asserting on trainer.processing_class and no helper at all.

Happy either way on timing: push the fixture in this PR, or merge the fix now and do it as a follow-up. Your call.

One more thing, passing a model instance leaves model_revision = None, so that path still loads the processor from the default branch. fine by me, config._commit_hash isn't worth it IMO, just mention it in the PR description?

[behroozazarkhalili]

The fix is correct. Two propagation gaps beyond the ones already listed, and the fixture the test discussion assumed does not exist is already on the Hub.

What I checked

RED/GREEN on all 11 new test ids. Base copy is the PR worktree with trl/ reset to 6630e17a and the new tests kept.

HEAD: 11 passed, 998 deselected, 2 warnings in 60.71s
BASE: 11 failed, 998 deselected, 2 warnings in 59.58s
      AssertionError: expected every processing_class `from_pretrained` call for
      'trl-internal-testing/tiny-Qwen2ForCausalLM-2.5' to use revision='main', got [None, None]

End-to-end reproduction of the issue, SFTTrainer(model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", args=SFTConfig(model_init_kwargs={"revision": "8913f5819566"})), then len(trainer.processing_class.chat_template):

BASE: CHAT_TEMPLATE_LEN 2507      # main's template, not the pinned revision's
HEAD: CHAT_TEMPLATE_LEN 2558      # the pinned revision's template

ruff check and ruff format --check at the CI-pinned 0.13.3 pass on all 23 changed files.

Also confirmed there is no processing_class auto-load anywhere in trl/scripts/ or trl/cli.py, so the seven CLI entry points are fixed purely through model_init_kwargs.

Findings

1. The chat-template fixture already exists, so the strong test is writable in this PR. trl-internal-testing/tiny-Qwen2ForCausalLM-2.5 has 18 commits and its chat template changed along the way: chat_template.jinja is 2558 bytes (sha1 90fd14a1) at revision 8913f5819566 and 2507 bytes (sha1 9755ed0c) at main. No new Hub repo needed. That is the output quoted above, and it fails on base for the reason the issue describes rather than on a missing kwarg. Two supporting points for dropping expected_revision="main": hardcoding revision="main" at sft_trainer.py:1002 passes the new SFT test (I ran the mutant), and the helper records the same load twice, seen_revisions is [None, None] on base because more than one of the four patched classes sees a single auto-load.

2. revision is not forwarded to the reward tokenizer, in three trainers this PR touches plus online_dpo. grpo_trainer.py:521-529 loads a string reward model with args.model_init_kwargs, so the reward model already comes from revision, while grpo_trainer.py:563 auto-loads its tokenizer without it. Same pair in rloo_trainer.py:483, sdpo_trainer.py:666 and online_dpo_trainer.py:241 (whose reward model is likewise loaded with args.model_init_kwargs at :214-221). This is not a separate scope question: grpo_config.py:437-439 already documents trust_remote_code as "Also applied to reward-model and reward-tokenizer loads", so the reward path is inside the forwarding contract these configs describe. Suggest forwarding the same model_init_kwargs.get("revision") there, in the same commit.

3. gold_trainer.py:831 and :868, teacher processor ignores teacher_model_revision. gold_trainer.py:962-964 does init_kwargs.setdefault("revision", args.teacher_model_revision) for the teacher weights, but the teacher processor loaded 130 lines earlier does not, and it is what becomes self._teacher_processor for the ULD loss. Two more rows for the table above, same shape as the iw_opd asymmetry.

4. reward_config.py:105, minor doc bug worth fixing in the same pass. The model_init_kwargs help says AutoModelForCausalLM.from_pretrained, but RewardTrainer loads with AutoModelForSequenceClassification (reward_trainer.py:393), and reward_trainer.py:255 states it correctly. Same two lines that need the AutoTokenizer mention added.

One consistency note on the diff itself: eight trainers set model_revision = None in the else branch, while a2po, sdft and sdpo pre-initialise it before the if. That is forced by those three having no else branch, so it reads fine, just flagging that the two idioms are deliberate.

Not reproduced

Nothing. Every claim above was executed or read at the cited line.

Unevaluable

Whether the auto-loaded processor and the model can disagree on anything other than the chat template at a pinned revision. Only the template differs between the two revisions I found on the fixture repo.

[kashif]

Updated from current main and addressed the review feedback in 7546ad6 and 172e0f1: completed the experimental propagation sweep, forwarded reward and GOLD teacher revisions, replaced the shared mock helper with the pinned chat-template regression test, corrected RewardConfig, and documented the revision-only scope plus model-instance behavior. make precommit passes, and the focused SFT regression test passes.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit 172e0f1. Configure here.}

[qgallouedec]

Test point addressed, thanks.

Good catch on the ULD one.

One thing left from my last comment: the description still doesn't mention that passing a model instance leaves model_revision = None, so that path keeps loading the processor from the default branch.
You've implemented it that way deliberately with the isinstance(model, str) guard, it just isn't written down. Add a line and we're good to merge.

Small nit, worth a comment next to assert len(...) == 2558 saying main is 2507. Otherwise it reads as an arbitrary number and the next person won't know what it's protecting.