ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

CI(Core): add compiler-cache coverage (synthetic invariants + real-class round-trip)

Browse source 
Adds two new strict-mode steps to the Core matrix to exercise the
dynamic file generation path in unsloth_zoo.compiler. Synthesized from
parallel design forks (cache_invariants + real-class + monkey-patch);
matrix expansion + monkey-patches stay as future PRs.

Step 1 -- "Compiler cache hygiene + source-rewriter invariants
(synthetic inputs)" -- 9 pytest cases on tiny synthetic source strings.
Covers higher_precision_softmax (basic + idempotent),
fix_rotary_embedding_dtype (no-op + active),
fix_attention_dtype_consistency (insert + idempotent),
convert_attention_masks_to_bool (rewrite + no-op),
create_new_function happy-path (versioning block / license header /
ast.parse / importlib re-import), and the UNSLOTH_COMPILE_OVERWRITE=0
forced-recompile-on-version-mismatch + matching-versions short-circuit
branches at compiler.py:947-963. Wall-time ~10-25s per cell.

Step 2 -- "Compiler real-class round-trip (llama / qwen3 / gemma3 +
SFT trainer)" -- runs unsloth_compile_transformers against actual
transformers modeling modules (llama, qwen3, gemma3) and TRL's
SFTTrainer. ast.parse + importlib + surface check on each generated
unsloth_compiled_cache/*.py. Includes a negative control test that
DISABLE=1 writes nothing. Hermetic per-pytest tempdir; skips legitimately
when transformers lacks a target model_type. Wall-time ~2-3 min per cell.

Both steps reuse tests/_zoo_aggressive_cuda_spoof.py and follow the
same auto-write-shim pattern as _zoo_apply_fused_lm_head_shim. The
job-level UNSLOTH_COMPILE_DISABLE=1 is popped inside the round-trip
shim so compilation actually fires there; restored on exit.

Plans at plans/compiler_cache_ci_fork_{a,b,c}.md (fork C's 3x3 matrix
expansion + NEEDS_PRECONDITION lift via monkey-patch are out of scope
for this PR but tracked there for follow-up).
This commit is contained in:
Daniel Han 2026-05-07 08:04:29 +00:00
commit 20e06ac52e
1 changed files with 323 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

323
.github/workflows/consolidated-tests-ci.yml vendored
View file

@ -599,6 +599,329 @@ jobs:
python -m pytest -q --tb=short tests/_tiled_mlp_check_shim.py -s
rm -f tests/_tiled_mlp_check_shim.py
- name: Compiler cache hygiene + source-rewriter invariants (synthetic inputs)
# Lightweight pipeline coverage for unsloth_zoo.compiler. Pure regex
# / tokenize / ast paths driven by tiny synthetic source strings:
# - higher_precision_softmax (basic + idempotent)
# - fix_rotary_embedding_dtype (no-op + active under
# UNSLOTH_FORCE_CUSTOM_DTYPE)
# - fix_attention_dtype_consistency (insert + idempotent)
# - convert_attention_masks_to_bool (rewrite + no-op)
# - create_new_function happy-path (versioning block, license
# header, AST parse, importlib re-import)
# - create_new_function **kwargs collision (exercises
# _rewrite_kwargs_param + _insert_kwargs_alias)
# - UNSLOTH_COMPILE_OVERWRITE=0 forced-recompile on transformers
# version mismatch (compiler.py:947-963)
# - matching short-circuit when versions are equal
# No real transformers modeling module is loaded; complements the
# heavier real-class round-trip step below. Wall-time ~10-25s.
run: |
set -euxo pipefail
cat > tests/_compiler_cache_invariants_shim.py <<'PY'
# Auto-generated by .github/workflows/consolidated-tests-ci.yml.
# Cache-hygiene + source-rewriter invariants for unsloth_zoo.compiler.
import sys, pathlib, os, ast, importlib, importlib.util, time
sys.path.insert(0, str(pathlib.Path(__file__).parent))
import _zoo_aggressive_cuda_spoof as _spoof
_spoof.apply()
import pytest
import torch # noqa: F401 (compiler.py imports torch at module load)
def _isolate_cache(tmp_path, monkeypatch):
"""Point UNSLOTH_COMPILE_LOCATION at tmp_path and reset module
globals. The compiler.py global is captured at module load
(line 75/179), so we delete + reimport per test."""
monkeypatch.setenv("UNSLOTH_COMPILE_LOCATION", str(tmp_path))
if "unsloth_zoo.compiler" in sys.modules:
del sys.modules["unsloth_zoo.compiler"]
import unsloth_zoo.compiler as compiler
compiler.UNSLOTH_COMPILE_LOCATION = str(tmp_path)
compiler.UNSLOTH_COMPILE_USE_TEMP = False
return compiler
def test_higher_precision_softmax_basic_and_idempotent(tmp_path, monkeypatch):
c = _isolate_cache(tmp_path, monkeypatch)
src = (
"y = nn.functional.softmax(x, dim=-1)\n"
"z = F.softmax(a, dim=1, dtype=torch.bfloat16)\n"
)
out = c.higher_precision_softmax(src)
assert "dtype = torch.float32).to(x.dtype)" in out
assert "dtype = torch.float32).to(a.dtype)" in out
assert c.higher_precision_softmax(out) == out
def test_fix_rotary_dtype_no_op_without_env(tmp_path, monkeypatch):
c = _isolate_cache(tmp_path, monkeypatch)
monkeypatch.delenv("UNSLOTH_FORCE_CUSTOM_DTYPE", raising=False)
src = "out = cos.to(dtype=x.dtype) + sin.to(dtype=x.dtype)\n"
assert c.fix_rotary_embedding_dtype(src) == src
def test_fix_rotary_dtype_active(tmp_path, monkeypatch):
c = _isolate_cache(tmp_path, monkeypatch)
monkeypatch.setenv(
"UNSLOTH_FORCE_CUSTOM_DTYPE",
"float16;torch.float32;torch.bfloat16;torch.float16;pass",
)
monkeypatch.setenv("UNSLOTH_FORCE_FLOAT32", "1")
src = "out = cos.to(dtype=x.dtype) + sin.to(dtype=x.dtype)\n"
out = c.fix_rotary_embedding_dtype(src)
# Active form rewrites cos.to / sin.to. Either the conditional
# form or the cast form is acceptable -- different transformers
# versions surface slightly different outputs from the rewriter.
assert "cos.to(dtype=x.dtype)" not in out
assert "sin.to(dtype=x.dtype)" not in out
def test_fix_attention_dtype_consistency_insert_then_idempotent(tmp_path, monkeypatch):
c = _isolate_cache(tmp_path, monkeypatch)
src = (
" query_states, key_states = apply_rotary_pos_emb("
"query_states, key_states, cos, sin)\n"
" attn = q @ k.T\n"
)
out = c.fix_attention_dtype_consistency(src)
assert out.count("value_states = value_states.to(query_states.dtype)") == 1
assert c.fix_attention_dtype_consistency(out) == out
def test_convert_attention_masks_to_bool_rewrites(tmp_path, monkeypatch):
c = _isolate_cache(tmp_path, monkeypatch)
src = (
"def make_mask(x):\n"
" out = torch.finfo(x.dtype).min * x\n"
" return out\n"
)
out = c.convert_attention_masks_to_bool("make_mask", src)
# Loose match: rewriter inserts a `!=torch.finfo(...).min` check
# somewhere on the return path. Tightening to an exact
# last-line match is brittle across transformers versions.
assert "!=torch.finfo" in out
def test_convert_attention_masks_to_bool_no_op(tmp_path, monkeypatch):
c = _isolate_cache(tmp_path, monkeypatch)
src = "def make_mask(x):\n return x\n"
assert c.convert_attention_masks_to_bool("make_mask", src) == src
def _versioning_lines(file_text):
"""Extract the four version strings from the versioning block."""
assert file_text.startswith('"""\n'), "missing opening triple-quote"
head = file_text.split("__UNSLOTH_VERSIONING__", 1)[0]
lines = [ln for ln in head.splitlines() if ln and ln != '"""']
return lines
def test_create_new_function_happy_path(tmp_path, monkeypatch):
c = _isolate_cache(tmp_path, monkeypatch)
src = "def f(x):\n return nn.functional.softmax(x, dim=-1)\n"
c.create_new_function(
name="f_happy", new_source=src, model_location="builtins",
functions=[], overwrite=True,
)
cached = tmp_path / "f_happy.py"
assert cached.exists()
text = cached.read_text(encoding="utf-8")
versions = _versioning_lines(text)
assert len(versions) == 4, versions
assert text.count(c._full_license_header) == 1
ast.parse(text)
spec = importlib.util.spec_from_file_location("f_happy_reimport", cached)
m2 = importlib.util.module_from_spec(spec)
spec.loader.exec_module(m2)
assert callable(m2.f)
import inspect as _inspect
# higher_precision_softmax should have promoted to float32.
assert "dtype = torch.float32" in _inspect.getsource(m2.f)
def test_create_new_function_overwrite_zero_recompiles_on_version_mismatch(
tmp_path, monkeypatch,
):
c = _isolate_cache(tmp_path, monkeypatch)
name = "vmismatch"
cached = tmp_path / f"{name}.py"
stub = (
'"""\n0.0.0\n0.0.0\n0.0.0-stub\n0.0.0\n__UNSLOTH_VERSIONING__\n"""\n'
+ c._full_license_header
+ "def vmismatch(x):\n return x\n"
)
cached.write_text(stub, encoding="utf-8")
monkeypatch.setenv("UNSLOTH_COMPILE_OVERWRITE", "0")
src = "def vmismatch(x):\n return x + 1\n"
c.create_new_function(
name=name, new_source=src, model_location="builtins",
functions=[], overwrite=False,
)
text = cached.read_text(encoding="utf-8")
assert "0.0.0-stub" not in text, (
"OVERWRITE=0 + transformers-version-mismatch did NOT recompile"
)
versions = _versioning_lines(text)
import importlib.metadata as _md
assert versions[2] == _md.version("transformers")
def test_create_new_function_overwrite_zero_short_circuits_when_versions_match(
tmp_path, monkeypatch,
):
c = _isolate_cache(tmp_path, monkeypatch)
name = "vmatch"
src = "def vmatch(x):\n return x\n"
c.create_new_function(
name=name, new_source=src, model_location="builtins",
functions=[], overwrite=True,
)
cached = tmp_path / f"{name}.py"
mtime_before = cached.stat().st_mtime_ns
time.sleep(0.05)
monkeypatch.setenv("UNSLOTH_COMPILE_OVERWRITE", "0")
c.create_new_function(
name=name, new_source=src, model_location="builtins",
functions=[], overwrite=False,
)
assert cached.stat().st_mtime_ns == mtime_before, (
"OVERWRITE=0 + matching versions should NOT rewrite the file"
)
PY
python -m pytest -q --tb=short tests/_compiler_cache_invariants_shim.py
rm -f tests/_compiler_cache_invariants_shim.py
- name: Compiler real-class round-trip (llama / qwen3 / gemma3 + SFT trainer)
# Heavier complementary path to the cache-hygiene step above.
# Calls `unsloth_compile_transformers(model_type=...)` against
# actual transformers modeling modules and `_patch_trl_rl_trainers`
# against TRL's SFTTrainer, then ast.parse / importlib-load /
# introspect the generated unsloth_compiled_cache/*.py files.
# Catches regex / source-rewriter drift across the matrix's
# (transformers, trl) combination -- the dominant failure mode of
# `unsloth_compile_transformers` after a transformers point release.
# Hermetic cache dir per pytest invocation; we override the
# job-level UNSLOTH_COMPILE_DISABLE=1 inside the shim so
# compilation actually runs here. Wall-time ~2-3 min.
run: |
set -euxo pipefail
cat > tests/_zoo_compiler_cache_shim.py <<'PY'
# Auto-generated by .github/workflows/consolidated-tests-ci.yml.
import os, sys, ast, pathlib, importlib.util, tempfile
_HERE = pathlib.Path(__file__).parent
sys.path.insert(0, str(_HERE))
import _zoo_aggressive_cuda_spoof as _spoof
_spoof.apply()
# Hermetic cache dir + force compile path BEFORE importing
# unsloth_zoo.compiler (its globals capture env at module load).
_CACHE = pathlib.Path(tempfile.mkdtemp(prefix="unsloth_cache_"))
os.environ["UNSLOTH_COMPILE_LOCATION"] = str(_CACHE)
os.environ["UNSLOTH_COMPILE_OVERWRITE"] = "1"
os.environ.pop("UNSLOTH_COMPILE_DISABLE", None)
import pytest
from unsloth_zoo.compiler import unsloth_compile_transformers
def _verify_file(path: pathlib.Path, must_expose):
assert path.exists(), f"compiler did not write {path}"
src = path.read_text(encoding="utf-8")
ast.parse(src, filename=str(path))
spec = importlib.util.spec_from_file_location(path.stem, path)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
for name in must_expose:
assert hasattr(mod, name), (
f"{path.name} missing expected attr {name!r}; "
f"found: {sorted(n for n in dir(mod) if not n.startswith('_'))[:25]}"
)
@pytest.mark.parametrize("model_type,rms_class", [
("llama", "LlamaRMSNorm"),
("qwen3", "Qwen3RMSNorm"),
("gemma3", "Gemma3RMSNorm"),
])
def test_compile_real_modeling_module(model_type, rms_class):
import importlib as _il
try:
_il.import_module(
f"transformers.models.{model_type}.modeling_{model_type}"
)
except ModuleNotFoundError:
pytest.skip(
f"transformers build lacks model_type={model_type}"
)
# fast_lora_forwards=False: the LoRA path expects PEFT + a real
# device for some torch.compile builds; skip it here, the
# source-emission path is what we want to verify.
unsloth_compile_transformers(
model_type=model_type, fast_lora_forwards=False,
)
modeling = _il.import_module(
f"transformers.models.{model_type}.modeling_{model_type}"
)
assert getattr(modeling, "__UNSLOTH_PATCHED__", False) is True
combined = _CACHE / f"unsloth_compiled_module_{model_type}.py"
_verify_file(combined, must_expose=[rms_class])
def test_compile_disable_writes_nothing():
"""Negative control: when UNSLOTH_COMPILE_DISABLE=1 the
compile path must early-return without producing new files."""
os.environ["UNSLOTH_COMPILE_DISABLE"] = "1"
try:
before = set(_CACHE.iterdir())
# Pick a model_type that still resolves on this transformers.
for mt in ("llama", "mistral", "qwen2"):
try:
import importlib as _il
_il.import_module(
f"transformers.models.{mt}.modeling_{mt}"
)
break
except ModuleNotFoundError:
continue
else:
pytest.skip("no probe model_type available")
unsloth_compile_transformers(
model_type=mt, fast_lora_forwards=False,
)
after = set(_CACHE.iterdir())
assert after == before, (
f"DISABLE=1 still wrote: {[p.name for p in after - before]}"
)
finally:
os.environ.pop("UNSLOTH_COMPILE_DISABLE", None)
def test_compile_sft_trainer_patch():
"""Round-trip TRL's SFTTrainer through the rl.py patch path
and verify the generated UnslothSFTTrainer.py."""
pytest.importorskip("trl")
try:
from unsloth.models.rl import _patch_trl_rl_trainers
except ImportError:
pytest.skip("unsloth.models.rl._patch_trl_rl_trainers absent")
try:
_patch_trl_rl_trainers("sft_trainer")
except Exception as e:
# TRL 1.x renames break the patch helper internally; we
# accept that here and skip rather than fail the cell.
pytest.skip(f"_patch_trl_rl_trainers raised: {type(e).__name__}: {e}")
sft = _CACHE / "UnslothSFTTrainer.py"
if not sft.exists():
pytest.skip(
"_patch_trl_rl_trainers ran but did not emit "
"UnslothSFTTrainer.py on this TRL version."
)
_verify_file(sft, must_expose=["UnslothSFTTrainer"])
PY
python -m pytest -q --tb=short tests/_zoo_compiler_cache_shim.py
rm -f tests/_zoo_compiler_cache_shim.py
- name: llama.cpp install via unsloth_zoo.llama_cpp + `llama-cli --help` smoke
# Exercise the canonical `unsloth_zoo.llama_cpp.install_llama_cpp`
# flow that GGUF export uses at runtime: clone ggml-org/llama.cpp