Skip to content

version-compat CI: fake CPU training runs for SFT/GRPO/DPO - #6965

Merged
danielhanchen merged 6 commits into
mainfrom
feat/cpu-fake-training-ci
Jul 8, 2026
Merged

danielhanchen merged 6 commits into
mainfrom
feat/cpu-fake-training-ci

Conversation

@danielhanchen

Copy link
Copy Markdown
Member

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 real train() 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:

  • Applies the aggressive CUDA spoof, then a small chain of CPU shims documented inline: eager torch.compile passthrough, dynamo suppression, device="cuda" tensor-alloc redirected to CPU, model.for_training/for_inference equivalents, and a couple of stream-capture stubs Adam probes.
  • Uses a dense (non-MoE) tiny model on purpose, since MoE routes through the CUDA-only grouped_gemm kernel.
  • Skips cleanly when torch/unsloth/trl are absent or the hub is unreachable, so it never turns an infra hiccup into a red build.

.github/workflows/version-compat-ci.yml: the grpo-fake-run job now runs this file alongside the patch-run test, against TRL latest (always) and TRL main (on push).

Notes

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +173 to +181
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,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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.

Suggested change
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,

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +189 to +197
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,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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.

Suggested change
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,

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in d2940f8 (tmp_path).

Comment on lines +209 to +217
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,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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.

Suggested change
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,

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in d2940f8 (tmp_path).

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread unsloth/models/rl.py
Comment on lines +1658 to +1661
peft_pattern = (
r"\s*elif is_peft_model\(model\) and args\.beta != 0\.0:"
r".*?"
r"ref_param\.data\.copy_\(param\.data\)"

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +576 to +577
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")

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

@danielhanchen
danielhanchen force-pushed the feat/cpu-fake-training-ci branch from d2940f8 to 0a4c311 Compare July 8, 2026 10:40
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Keep them coming!

Reviewed commit: 202ccf0b10

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

danielhanchen and others added 6 commits July 8, 2026 11:05
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).
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.
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.
@danielhanchen
danielhanchen force-pushed the feat/cpu-fake-training-ci branch from 202ccf0 to d901e5d Compare July 8, 2026 11:06
@danielhanchen
danielhanchen merged commit 0e1ed88 into main Jul 8, 2026
20 of 24 checks passed
@danielhanchen
danielhanchen deleted the feat/cpu-fake-training-ci branch July 8, 2026 11:06
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.

1 participant