version-compat CI: fake CPU training runs for SFT/GRPO/DPO - #6965
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces compatibility updates for TRL >= 1.7.0, specifically handling the new 3-tuple return contract for per-token log probabilities, neutralizing the updated QLoRA bf16 cast, and adding a fail-fast for unsupported MoE router auxiliary loss. It also adds comprehensive version compatibility tests, including fake CPU training runs and source-patching validation. The feedback suggests utilizing pytest's built-in tmp_path fixture in the CPU training tests instead of hardcoded relative paths to ensure proper cleanup of temporary directories.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def test_sft_trains_on_cpu(): | ||
| from datasets import Dataset | ||
| from trl import SFTConfig, SFTTrainer | ||
|
|
||
| assert SFTTrainer.__name__ == "UnslothSFTTrainer", "SFT patch did not apply" | ||
| model, tok = _load_plain() | ||
| ds = Dataset.from_list([{"text": "The quick brown fox jumps over the lazy dog."}] * 8) | ||
| cfg = SFTConfig( | ||
| output_dir = "temp/ci_sft", per_device_train_batch_size = 2, max_steps = 2, |
There was a problem hiding this comment.
Using a hardcoded relative path like "temp/ci_sft" for output_dir leaves persistent directories in the workspace after tests run. It is highly recommended to use pytest's built-in tmp_path fixture to automatically manage and clean up temporary directories.
| def test_sft_trains_on_cpu(): | |
| from datasets import Dataset | |
| from trl import SFTConfig, SFTTrainer | |
| assert SFTTrainer.__name__ == "UnslothSFTTrainer", "SFT patch did not apply" | |
| model, tok = _load_plain() | |
| ds = Dataset.from_list([{"text": "The quick brown fox jumps over the lazy dog."}] * 8) | |
| cfg = SFTConfig( | |
| output_dir = "temp/ci_sft", per_device_train_batch_size = 2, max_steps = 2, | |
| def test_sft_trains_on_cpu(tmp_path): | |
| from datasets import Dataset | |
| from trl import SFTConfig, SFTTrainer | |
| assert SFTTrainer.__name__ == "UnslothSFTTrainer", "SFT patch did not apply" | |
| model, tok = _load_plain() | |
| ds = Dataset.from_list([{"text": "The quick brown fox jumps over the lazy dog."}] * 8) | |
| cfg = SFTConfig( | |
| output_dir = str(tmp_path / "ci_sft"), per_device_train_batch_size = 2, max_steps = 2, |
There was a problem hiding this comment.
Fixed in d2940f8: switched each trainer's output_dir to pytest tmp_path, so a local run no longer leaves untracked temp/ci_* dirs in the repo tree and the tests are CWD-independent.
| def test_grpo_trains_on_cpu(): | ||
| from datasets import Dataset | ||
| from trl import GRPOConfig, GRPOTrainer | ||
|
|
||
| assert GRPOTrainer.__name__ == "UnslothGRPOTrainer", "GRPO patch did not apply" | ||
| model, tok = _load_plain() | ||
| ds = Dataset.from_list([{"prompt": "hi there"}] * 4) | ||
| cfg = GRPOConfig( | ||
| output_dir = "temp/ci_grpo", per_device_train_batch_size = 2, num_generations = 2, |
There was a problem hiding this comment.
Using a hardcoded relative path like "temp/ci_grpo" for output_dir leaves persistent directories in the workspace after tests run. It is highly recommended to use pytest's built-in tmp_path fixture to automatically manage and clean up temporary directories.
| def test_grpo_trains_on_cpu(): | |
| from datasets import Dataset | |
| from trl import GRPOConfig, GRPOTrainer | |
| assert GRPOTrainer.__name__ == "UnslothGRPOTrainer", "GRPO patch did not apply" | |
| model, tok = _load_plain() | |
| ds = Dataset.from_list([{"prompt": "hi there"}] * 4) | |
| cfg = GRPOConfig( | |
| output_dir = "temp/ci_grpo", per_device_train_batch_size = 2, num_generations = 2, | |
| def test_grpo_trains_on_cpu(tmp_path): | |
| from datasets import Dataset | |
| from trl import GRPOConfig, GRPOTrainer | |
| assert GRPOTrainer.__name__ == "UnslothGRPOTrainer", "GRPO patch did not apply" | |
| model, tok = _load_plain() | |
| ds = Dataset.from_list([{"prompt": "hi there"}] * 4) | |
| cfg = GRPOConfig( | |
| output_dir = str(tmp_path / "ci_grpo"), per_device_train_batch_size = 2, num_generations = 2, |
| def test_dpo_trains_on_cpu(): | ||
| from datasets import Dataset | ||
| from trl import DPOConfig, DPOTrainer | ||
|
|
||
| assert DPOTrainer.__name__ == "UnslothDPOTrainer", "DPO patch did not apply" | ||
| model, tok = _load_plain() | ||
| ds = Dataset.from_list([{"prompt": "Hi", "chosen": " hello friend", "rejected": " go away"}] * 8) | ||
| cfg = DPOConfig( | ||
| output_dir = "temp/ci_dpo", per_device_train_batch_size = 2, max_steps = 2, |
There was a problem hiding this comment.
Using a hardcoded relative path like "temp/ci_dpo" for output_dir leaves persistent directories in the workspace after tests run. It is highly recommended to use pytest's built-in tmp_path fixture to automatically manage and clean up temporary directories.
| def test_dpo_trains_on_cpu(): | |
| from datasets import Dataset | |
| from trl import DPOConfig, DPOTrainer | |
| assert DPOTrainer.__name__ == "UnslothDPOTrainer", "DPO patch did not apply" | |
| model, tok = _load_plain() | |
| ds = Dataset.from_list([{"prompt": "Hi", "chosen": " hello friend", "rejected": " go away"}] * 8) | |
| cfg = DPOConfig( | |
| output_dir = "temp/ci_dpo", per_device_train_batch_size = 2, max_steps = 2, | |
| def test_dpo_trains_on_cpu(tmp_path): | |
| from datasets import Dataset | |
| from trl import DPOConfig, DPOTrainer | |
| assert DPOTrainer.__name__ == "UnslothDPOTrainer", "DPO patch did not apply" | |
| model, tok = _load_plain() | |
| ds = Dataset.from_list([{"prompt": "Hi", "chosen": " hello friend", "rejected": " go away"}] * 8) | |
| cfg = DPOConfig( | |
| output_dir = str(tmp_path / "ci_dpo"), per_device_train_batch_size = 2, max_steps = 2, |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 105848dc0d
ℹ️ 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".
| peft_pattern = ( | ||
| r"\s*elif is_peft_model\(model\) and args\.beta != 0\.0:" | ||
| r".*?" | ||
| r"ref_param\.data\.copy_\(param\.data\)" |
There was a problem hiding this comment.
Handle DPO's TRL 1.7 PEFT ref branch
For TRL 1.7.x this new removal only matches GRPO's elif is_peft_model(model) and args.beta != 0.0: branch. DPO changed the analogous ref-adapter path to elif is_peft_model(model) and ref_model is None, and the generic remover later in this file still only matches the older if is_peft_available() ... args.beta shape, so UnslothDPOTrainer with a pre-wrapped PEFT model will still create a ref adapter and compute reference logits from the copied adapter instead of the base model. Please add the DPO 1.7 branch to the removal or cover it with a PEFT DPO canary.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The observation is accurate but this is not a reachable correctness bug for this PR, which only adds tests and CI. Confirmed against TRL 1.7.1: DPO's ref-adapter branch is elif is_peft_model(model) and ref_model is None, and the entire PEFT-removal block in rl.py is nested under if RLTrainer_name == "GRPOTrainer", so DPO's block is never matched, in any TRL version, and this PR does not change that. The surviving behavior is also intended rather than wrong: for a pre-wrapped PEFT model TRL deliberately snapshots the initial adapter as the reference so DPO regularizes against the pre-DPO policy, and for a fresh zero-init LoRA that snapshot is logit-identical to the base model. Forcing a base-model reference (the suggested change) would alter DPO reference semantics and could regress continued-DPO runs, so it belongs in a dedicated rl.py PR after a semantics decision, not here. test_dpo_trains_on_cpu uses a plain non-PEFT model, so the branch is not even reached at runtime.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d2940f8437
ℹ️ 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 os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": | ||
| self._autocast_dtype = torch.float16 | ||
|
|
||
| compute_aux_loss = kwargs.get("compute_aux_loss", None) |
There was a problem hiding this comment.
Forward the new TRL 1.7 VLM kwargs
When TRL >= 1.7 calls this replacement for VLM batches, it can pass spatial_shapes, num_tiles, or image_position_ids in addition to compute_aux_loss (for example LFM2-VL/image-position processors). This compatibility path only reads the new aux kwarg and then chunks/forwards the older image_grid_thw/pixel_attention_mask fields, so those batches drop required vision metadata before unwrapped_model(...) and either mis-slice pixel_values or call the model without required args. Please include, slice, and forward the new VLM kwargs when enabling the 1.7 arity path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Accurate, and intentionally out of scope here: this PR only adds tests and CI. The new TRL 1.7 VLM kwargs (spatial_shapes / num_tiles / image_position_ids) are unreachable for text GRPO and are tracked in issue #6960, which was explicitly deferred from the TRL >= 1.7 crash/arity fix. It needs a GPU VLM parity run, so it belongs there rather than in this CI PR.
| return lambda fn: fn | ||
|
|
||
|
|
||
| torch.compile = _eager_compile |
There was a problem hiding this comment.
Move the compile shim before conftest imports Unsloth
In this new pytest job, tests/conftest.py is loaded before this module and it imports unsloth when torch is installed. If that import succeeds, the TRL trainers are generated before this assignment replaces torch.compile, so the GRPO hot functions can already be real Dynamo-compiled and the CPU fake-train can still hit CUDA/accelerator probes despite the shim. Set the env/compile shim in the workflow or conftest before the early unsloth import, or disable that import for this job.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The ordering observation is correct (conftest.py does import unsloth before this module's shim), so I added TORCHDYNAMO_DISABLE / TORCH_COMPILE_DISABLE to the fake-run step env, which is set before any import. To be precise though, the predicted failure did not actually occur before this: the workflow already sets UNSLOTH_COMPILE_DISABLE=1 in the step env before conftest imports unsloth, and the autouse fixture flips torch._dynamo.config.disable=True, a call-time kill switch that forces eager even on functions decorated earlier with the real torch.compile (verified: a function decorated pre-flag runs eager post-flag, grads intact). The suite passes on a real CPU-build runner. The new env vars are defense in depth.
| # and explicit opt-in is rejected at trainer init, so this is always None (kept in the | ||
| # return for TRL >= 1.7.0's 3-tuple contract). | ||
| aux_loss = None | ||
| return logprobs.detach(), entropies, aux_loss # logps, entropies, aux_loss |
There was a problem hiding this comment.
Return trainable logps and entropy for TRL 1.7
On TRL >= 1.7, GRPOTrainer uses its native _compute_loss path (compute_loss is inherited, so the existing function_name == "compute_loss" replacement is not installed) and unpacks this 3-tuple for the actual policy loss and entropy metrics. This return still detaches logprobs and leaves entropies as None, so latest-TRL GRPO training either crashes when _compute_loss logs entropy or produces a loss disconnected from model parameters. The 1.7 path needs differentiable policy logprobs and real entropy values, or _compute_loss must be patched to the Unsloth accumulated-loss path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Verified empirically on TRL 1.7.1 and this does not occur. TRL 1.7.x defines compute_loss in the GRPOTrainer class body (it is not inherited), so rl.py's def-name enumeration installs the Unsloth compute_loss replacement; training_step calls that (the Unsloth accumulated-loss path), and TRL's native _compute_loss is never reached. The detached logprobs from _get_per_token_logps_and_entropies are consumed only at the old_per_token_logps / ref_per_token_logps call sites (unpacked as _, _), where detaching is correct; entropies there is discarded. A real GRPO train step on 1.7.1 gives loss.requires_grad=True, loss.grad_fn=UnslothEfficientGRPOBackward, and non-zero grads on all trainable params (embed_tokens, q/k/v/o_proj, gate_proj, ...). So the loss is connected and nothing logs the None entropy.
| if not _tag_ge(tag, "1.7.0"): | ||
| pytest.skip(f"{tag}: pre-1.7.0 PEFT block handled by rl.py's 0.26/0.27 branches") |
There was a problem hiding this comment.
Cover the v1.5/v1.6 PEFT block
This skip now excludes the newly added v1.5/v1.6 tags, but those releases already use the same elif is_peft_model(model) and args.beta != 0.0: ref-adapter block as v1.7. Because rl.py's pre-1.7 regex still looks for the older if is_peft_available() and is_peft_model(...) form, the transform no-ops for those tags while this matrix stays green. Please run this contract for the versions that contain the elif block, not only >=1.7.0.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, fixed. Confirmed against the TRL sources that the elif is_peft_model(model) and args.beta != 0.0: ref-adapter block was introduced in TRL 1.4.0 (not 1.7.0) and is unchanged through 1.7.x, while rl.py gated its removal at >= 1.7.0, so for 1.4 <= TRL < 1.7 the transform fell through to the 0.27 branch (older if is_peft_available()... form) and no-oped. Lowered the removal gate to Version(1.4.0) and kept the 1.7.0-only router aux-loss fail-fast nested, and widened this contract test to run from 1.4.0 so v1.4.0/v1.5.0/v1.5.1/v1.6.0 are actually exercised (they pass).
d2940f8 to
0a4c311
Compare
|
@codex review |
|
Codex Review: Didn't find any major issues. Keep them coming! 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". |
Adds a runtime layer on top of the patch-run canary: actually runs trainer.train() for a couple of steps on a CPU-only runner under the CUDA spoof, wrapping a plain tiny HF model in the Unsloth-patched trainer. Exercises the real train() loop (collation, generation, the injected _get_per_token_logps_and_entropies, loss, backward, optimizer) so a TRL or transformers change that breaks the loop at runtime -- not just the source structure -- surfaces here. No GPU, no meaningful numerics. Needs a chain of small CPU shims (eager torch.compile, dynamo suppress, cuda tensor-alloc redirect to CPU, model.for_training/for_inference equivalents) documented inline. Does not exercise Unsloth's Triton/GPU kernels (CPU can't).
for more information, see https://pre-commit.ci
On a real CPU-build torch runner (GitHub CI) two things bit that a CUDA-build torch with GPUs hidden masked locally: - The default optimizer is adamw_8bit (bitsandbytes), whose is_on_gpu() check dies on CPU tensors. Force optim=adamw_torch in all three configs. - import unsloth reinstalls the real torch.compile over the eager passthrough, so the GRPO hot path (chunked_selective_log_softmax) actually compiles and inductor picks the spoofed CUDA device, crashing on device props (gcnArchName). Re-apply the eager passthrough after import and flip torch._dynamo.config.disable so every @torch.compile runs eager at call time.
for more information, see https://pre-commit.ci
Use pytest's tmp_path for each trainer's output_dir instead of a hardcoded relative temp/ci_* path, so a local pytest run does not leave untracked dirs in the repo tree and the tests are CWD-independent.
Set TORCHDYNAMO_DISABLE / TORCH_COMPILE_DISABLE in the fake-run step env so dynamo/inductor is off before conftest.py's early import unsloth, not only via the per-test runtime shim. Defense in depth on the GPU-less runner: the GRPO hot path never compiles regardless of when its functions were decorated.
202ccf0 to
d901e5d
Compare
What
Adds a runtime layer to the version-compat CI that actually runs
trainer.train()for a couple of steps on a CPU-only GitHub runner, for SFT, GRPO and DPO, wrapping a plain tiny HF model in the Unsloth-patched trainer under the CUDA spoof.This complements the existing patch-run canary (
test_trl_grpo_fake_run.py), which only generates and inspects the transformed trainer source. The new test drives the realtrain()loop, data collation, generation (GRPO), the injected_get_per_token_logps_and_entropies, loss, backward and optimizer, so a TRL or transformers change that breaks the loop at runtime (not just the source structure) surfaces automatically.No GPU, no meaningful numerics; it validates the trainer-transform and orchestration layer with a standard forward, not the Triton/GPU kernels (those cannot run on CPU).
How
tests/version_compat/test_trl_fake_train_cpu.py:torch.compilepassthrough, dynamo suppression,device="cuda"tensor-alloc redirected to CPU,model.for_training/for_inferenceequivalents, and a couple of stream-capture stubs Adam probes..github/workflows/version-compat-ci.yml: thegrpo-fake-runjob now runs this file alongside the patch-run test, against TRL latest (always) and TRL main (on push).Notes