version-compat CI: fake CPU training runs for SFT/GRPO/DPO (#6965)
* version-compat CI: fake CPU training runs for SFT/GRPO/DPO 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). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cpu fake-train: force adamw_torch + disable dynamo for CPU runner 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cpu fake-train: write checkpoints under pytest tmp_path 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. * version-compat CI: disable dynamo at process level for the fake-run job 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. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
07c8bbbf5a
commit
0e1ed88bb8
2 changed files with 294 additions and 0 deletions
9
.github/workflows/version-compat-ci.yml
vendored
9
.github/workflows/version-compat-ci.yml
vendored
|
|
@ -339,12 +339,18 @@ jobs:
|
|||
env:
|
||||
UNSLOTH_IS_PRESENT: '1'
|
||||
UNSLOTH_COMPILE_DISABLE: '1'
|
||||
# Disable dynamo/inductor at the process level, before conftest.py's early
|
||||
# `import unsloth`, so the GRPO hot path never compiles on the GPU-less runner
|
||||
# (defense in depth; the CPU fake-train also flips this at runtime).
|
||||
TORCHDYNAMO_DISABLE: '1'
|
||||
TORCH_COMPILE_DISABLE: '1'
|
||||
PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python
|
||||
run: |
|
||||
cd unsloth
|
||||
python -c "import trl; print('Resolved TRL', trl.__version__)"
|
||||
PYTHONPATH=. python -m pytest \
|
||||
tests/version_compat/test_trl_grpo_fake_run.py \
|
||||
tests/version_compat/test_trl_fake_train_cpu.py \
|
||||
-v --tb=short
|
||||
# `main` is scheduled/dispatch-only so PR jobs stay fast and a bleeding-edge
|
||||
# TRL break does not red every PR. github.event_name is valid in a step if.
|
||||
|
|
@ -353,6 +359,8 @@ jobs:
|
|||
env:
|
||||
UNSLOTH_IS_PRESENT: '1'
|
||||
UNSLOTH_COMPILE_DISABLE: '1'
|
||||
TORCHDYNAMO_DISABLE: '1'
|
||||
TORCH_COMPILE_DISABLE: '1'
|
||||
PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python
|
||||
run: |
|
||||
pip install --upgrade "git+https://github.com/huggingface/trl"
|
||||
|
|
@ -360,6 +368,7 @@ jobs:
|
|||
python -c "import trl; print('Resolved TRL', trl.__version__)"
|
||||
PYTHONPATH=. python -m pytest \
|
||||
tests/version_compat/test_trl_grpo_fake_run.py \
|
||||
tests/version_compat/test_trl_fake_train_cpu.py \
|
||||
-v --tb=short
|
||||
|
||||
# Daily-only: same suites but with --strict on importable upstream
|
||||
|
|
|
|||
285
tests/version_compat/test_trl_fake_train_cpu.py
Normal file
285
tests/version_compat/test_trl_fake_train_cpu.py
Normal file
|
|
@ -0,0 +1,285 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team.
|
||||
"""Fake CPU training runs for the Unsloth-patched SFT / GRPO / DPO trainers.
|
||||
|
||||
The patch-run canary (test_trl_grpo_fake_run.py) only compiles + inspects the
|
||||
generated trainer source. This goes one layer deeper: it actually runs
|
||||
`trainer.train()` for a couple of steps on a CPU-only runner, under the CUDA
|
||||
spoof, wrapping a plain (tiny, random-weight) HF model in the Unsloth-patched
|
||||
trainer. That exercises the real train() loop at runtime -- data collation,
|
||||
generation (GRPO), the injected `_get_per_token_logps_and_entropies`, loss,
|
||||
backward, optimizer -- so a TRL or transformers change that breaks the loop
|
||||
(not just the source structure) surfaces here. No GPU, no meaningful numerics.
|
||||
|
||||
What it does NOT cover: Unsloth's Triton/GPU-optimized model kernels (the
|
||||
FastLanguageModel fast path) cannot run on CPU, so this validates the
|
||||
trainer-transform + orchestration layer with a standard forward, not the
|
||||
optimized kernels.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
# CPU-only: no torch.compile / dynamo (it reaches into the CUDA accelerator), no
|
||||
# Unsloth kernel compile, no mixed precision. Must be set before torch/unsloth.
|
||||
os.environ.setdefault("UNSLOTH_COMPILE_DISABLE", "1")
|
||||
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
|
||||
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
|
||||
os.environ.setdefault("ACCELERATE_MIXED_PRECISION", "no")
|
||||
|
||||
import importlib
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
# torch is needed for everything below (daily-fresh-fetch collects this dir with
|
||||
# only pytest installed); skip the whole module cleanly when it is absent.
|
||||
if importlib.util.find_spec("torch") is None:
|
||||
pytest.skip(
|
||||
"torch not installed; fake CPU train needs the real runtime", allow_module_level = True
|
||||
)
|
||||
|
||||
# Apply the CUDA spoof before any unsloth-touching import.
|
||||
_SPOOF_DIR = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(_SPOOF_DIR))
|
||||
import _zoo_aggressive_cuda_spoof as _spoof # noqa: E402
|
||||
|
||||
_spoof.apply()
|
||||
|
||||
import torch # noqa: E402
|
||||
|
||||
|
||||
# The generated GRPO trainer hard-decorates hot functions with @torch.compile,
|
||||
# which dynamo processes even under the disable env vars, reaching into
|
||||
# torch.accelerator (real CUDA) on a GPU-less box. Make torch.compile an eager
|
||||
# passthrough before unsloth generates/imports the trainer -- same logic, no
|
||||
# dynamo. (An eager CPU run is exactly what we want here.)
|
||||
def _eager_compile(
|
||||
model = None,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
if callable(model):
|
||||
return model
|
||||
return lambda fn: fn
|
||||
|
||||
|
||||
torch.compile = _eager_compile
|
||||
|
||||
# Belt-and-suspenders: if any @torch.compile still routes through dynamo, let it
|
||||
# fall back to eager instead of crashing, and stop its stream-capture probe from
|
||||
# reaching torch.accelerator -> real CUDA on a GPU-less box.
|
||||
try:
|
||||
import torch._dynamo # noqa: E402
|
||||
torch._dynamo.config.suppress_errors = True
|
||||
except Exception:
|
||||
pass
|
||||
if hasattr(torch, "accelerator"):
|
||||
torch.accelerator.is_available = lambda *a, **k: False
|
||||
|
||||
|
||||
# Redirect any `device="cuda"` tensor allocation / `.to("cuda")` / `.cuda()` to
|
||||
# CPU. The aggressive spoof deliberately keeps real allocators, but a fake CPU
|
||||
# train needs cuda-targeted ops (e.g. inductor's init_gpu_context does
|
||||
# `torch.empty(1, device="cuda")`) to land on CPU instead of erroring.
|
||||
def _is_cuda_dev(d):
|
||||
try:
|
||||
return d is not None and torch.device(d).type == "cuda"
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
for _name in (
|
||||
"empty",
|
||||
"zeros",
|
||||
"ones",
|
||||
"full",
|
||||
"tensor",
|
||||
"arange",
|
||||
"randn",
|
||||
"rand",
|
||||
"randint",
|
||||
"empty_like",
|
||||
"zeros_like",
|
||||
"ones_like",
|
||||
):
|
||||
_orig = getattr(torch, _name, None)
|
||||
if _orig is None:
|
||||
continue
|
||||
|
||||
def _redir(
|
||||
*args,
|
||||
_orig = _orig,
|
||||
**kwargs,
|
||||
):
|
||||
if _is_cuda_dev(kwargs.get("device")):
|
||||
kwargs["device"] = "cpu"
|
||||
return _orig(*args, **kwargs)
|
||||
|
||||
setattr(torch, _name, _redir)
|
||||
|
||||
_orig_to = torch.Tensor.to
|
||||
|
||||
|
||||
def _to_cpu(self, *args, **kwargs):
|
||||
args = tuple("cpu" if _is_cuda_dev(a) else a for a in args)
|
||||
if _is_cuda_dev(kwargs.get("device")):
|
||||
kwargs["device"] = "cpu"
|
||||
return _orig_to(self, *args, **kwargs)
|
||||
|
||||
|
||||
torch.Tensor.to = _to_cpu
|
||||
torch.Tensor.cuda = lambda self, *a, **k: self
|
||||
|
||||
# Extra CUDA stubs the aggressive spoof lacks, needed to walk a real train():
|
||||
# Adam's _cuda_graph_capture_health_check() probes stream capture.
|
||||
torch.cuda.is_current_stream_capturing = lambda *a, **k: False
|
||||
try:
|
||||
import torch.cuda.graphs as _cg # noqa: E402
|
||||
_cg._cuda_isCurrentStreamCapturing = lambda *a, **k: False
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# A broken libmlx.so in the shared site-packages crashes transformers' Mac-only
|
||||
# is_mlx_array probe on Linux; disable it.
|
||||
try:
|
||||
import transformers.utils.generic as _g # noqa: E402
|
||||
_g._is_mlx_available = False
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# Dense (non-MoE) tiny model on purpose: MoE models route through Unsloth's
|
||||
# grouped_gemm Triton kernel, which is CUDA-only and cannot run on a CPU runner.
|
||||
_MODEL = "hf-internal-testing/tiny-random-LlamaForCausalLM"
|
||||
|
||||
|
||||
def _load_plain():
|
||||
"""Tiny plain HF model + tokenizer on CPU. Skips (not fails) if the model
|
||||
cannot be fetched -- that is a network/hub issue, not an unsloth regression."""
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
try:
|
||||
tok = AutoTokenizer.from_pretrained(_MODEL)
|
||||
model = AutoModelForCausalLM.from_pretrained(_MODEL, dtype = torch.float32)
|
||||
except OSError as e: # hub unreachable / model missing
|
||||
pytest.skip(f"could not fetch {_MODEL} (network/hub): {str(e)[:150]}")
|
||||
if tok.pad_token is None:
|
||||
tok.pad_token = tok.eos_token
|
||||
# Unsloth's GRPO path calls model.for_training()/for_inference() (added by
|
||||
# FastLanguageModel). A plain HF model lacks them; supply minimal train/eval
|
||||
# equivalents so the loop proceeds without the optimized wrapper.
|
||||
if not hasattr(model, "for_training"):
|
||||
model.for_training = lambda *a, **k: model.train()
|
||||
if not hasattr(model, "for_inference"):
|
||||
model.for_inference = lambda *a, **k: model.eval()
|
||||
return model.to("cpu"), tok
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _require_stack():
|
||||
global torch # the `import torch._dynamo` below would otherwise shadow it as local
|
||||
if importlib.util.find_spec("unsloth") is None or importlib.util.find_spec("trl") is None:
|
||||
pytest.skip("unsloth or trl not installed")
|
||||
# A real import failure is a regression we want to surface, so do not guard it.
|
||||
import unsloth # noqa: F401 -- patches TRL trainers to the Unsloth variants
|
||||
|
||||
# `import unsloth` reinstalls the real torch.compile (overwriting the eager
|
||||
# passthrough set at module load), so the GRPO hot path (chunked_selective_
|
||||
# log_softmax) would really compile -- and inductor picks the spoofed CUDA
|
||||
# device, crashing on device props (`gcnArchName`). Re-apply the eager
|
||||
# passthrough and flip dynamo's call-time kill switch so every @torch.compile
|
||||
# runs eager regardless of when it was decorated. CPU eager is what we want.
|
||||
torch.compile = _eager_compile
|
||||
try:
|
||||
import torch._dynamo # noqa: E402
|
||||
torch._dynamo.config.disable = True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
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,
|
||||
logging_steps = 1,
|
||||
report_to = "none",
|
||||
save_strategy = "no",
|
||||
use_cpu = True,
|
||||
max_length = None,
|
||||
padding_free = False,
|
||||
dataset_text_field = "text",
|
||||
fp16 = False,
|
||||
bf16 = False,
|
||||
optim = "adamw_torch",
|
||||
)
|
||||
SFTTrainer(model = model, processing_class = tok, args = cfg, train_dataset = ds).train()
|
||||
|
||||
|
||||
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,
|
||||
max_steps = 2,
|
||||
max_completion_length = 8,
|
||||
logging_steps = 1,
|
||||
report_to = "none",
|
||||
temperature = 1.0,
|
||||
beta = 0.0,
|
||||
save_strategy = "no",
|
||||
use_cpu = True,
|
||||
use_vllm = False,
|
||||
fp16 = False,
|
||||
bf16 = False,
|
||||
optim = "adamw_torch",
|
||||
)
|
||||
GRPOTrainer(
|
||||
model = model,
|
||||
processing_class = tok,
|
||||
reward_funcs = [lambda completions, **k: [float(len(c)) for c in completions]],
|
||||
args = cfg,
|
||||
train_dataset = ds,
|
||||
).train()
|
||||
|
||||
|
||||
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,
|
||||
logging_steps = 1,
|
||||
report_to = "none",
|
||||
save_strategy = "no",
|
||||
use_cpu = True,
|
||||
beta = 0.1,
|
||||
fp16 = False,
|
||||
bf16 = False,
|
||||
optim = "adamw_torch",
|
||||
)
|
||||
DPOTrainer(model = model, processing_class = tok, args = cfg, train_dataset = ds).train()
|
||||
Loading…
Add table
Add a link
Reference in a new issue