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

Merge branch 'main' into studio/api-keys

Browse source
This commit is contained in:
Roland Tannous 2026-04-10 22:14:33 +04:00 • committed by GitHub
commit 8302af06a9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 253 additions and 138 deletions
Download patch file Download diff file Expand all files Collapse all files

36
install.sh
View file

@ -94,6 +94,36 @@ run_install_cmd() {
return $_rc
}
# Install bitsandbytes on AMD ROCm hosts. Uses the continuous-release_main
# wheel for the ROCm 4-bit GEMV fix (bnb PR #1887, post-0.49.2); bnb <= 0.49.2
# NaNs at decode shape on every AMD GPU. Falls back to PyPI >=0.49.1 if the
# pre-release URL is unreachable. Drop the pin once bnb 0.50+ ships on PyPI.
_install_bnb_rocm() {
_label="$1"
_venv_py="$2"
case "$_ARCH" in
x86_64|amd64)
_bnb_whl_url="https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl"
;;
aarch64|arm64)
_bnb_whl_url="https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl"
;;
*)
_bnb_whl_url=""
;;
esac
if [ -n "$_bnb_whl_url" ]; then
substep "installing bitsandbytes for AMD ROCm (pre-release, PR #1887)..."
if run_install_cmd "$_label (pre-release)" uv pip install --python "$_venv_py" \
--force-reinstall --no-cache-dir --no-deps "$_bnb_whl_url"; then
return 0
fi
substep "[WARN] bnb pre-release unreachable; falling back to PyPI (4-bit decode broken on ROCm)" "$C_WARN"
fi
run_install_cmd "$_label (pypi fallback)" uv pip install --python "$_venv_py" \
--force-reinstall --no-cache-dir --no-deps "bitsandbytes>=0.49.1"
}
if [ "$_next_is_package" = true ]; then
echo "❌ ERROR: --package requires an argument." >&2
exit 1
@ -1296,8 +1326,7 @@ if [ "$_MIGRATED" = true ]; then
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*)
substep "installing bitsandbytes for AMD ROCm..."
run_install_cmd "install bitsandbytes (AMD)" uv pip install --python "$_VENV_PY" --force-reinstall --no-cache-dir --no-deps "bitsandbytes>=0.49.1"
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
# Repair ROCm torch if overwritten during migrated install
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
if [ -z "$_has_hip" ]; then
@ -1437,8 +1466,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*)
substep "installing bitsandbytes for AMD ROCm..."
run_install_cmd "install bitsandbytes (AMD)" uv pip install --python "$_VENV_PY" --force-reinstall --no-cache-dir --no-deps "bitsandbytes>=0.49.1"
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
;;
esac
fi

132
studio/backend/core/inference/inference.py
View file

@ -5,23 +5,8 @@
Core inference backend - streamlined
"""
# On AMD ROCm, Unsloth's global monkey-patching of transformers model classes
# (LlamaRotaryEmbedding, attention modules, etc.) causes HIP kernel crashes
# (_assert_async_cuda_kernel -> HSA_STATUS_ERROR_EXCEPTION) during inference.
# Training works because it uses different code paths, but generation triggers
# the incompatible patched kernels. Skip the Unsloth import entirely on ROCm
# so transformers classes stay unmodified; the GGUF inference path (llama-server)
# is unaffected since it never imports these Python model classes.
_IS_ROCM_ENV = getattr(__import__("torch").version, "hip", None) is not None
if _IS_ROCM_ENV:
FastLanguageModel = None # Loaded on-demand only on NVIDIA
FastVisionModel = None
get_chat_template = None
else:
from unsloth import FastLanguageModel, FastVisionModel
from unsloth.chat_templates import get_chat_template
from unsloth import FastLanguageModel, FastVisionModel
from unsloth.chat_templates import get_chat_template
from transformers import TextStreamer
from peft import PeftModel, PeftModelForCausalLM
@ -41,7 +26,6 @@ from utils.hardware import (
raise_if_offloaded,
get_visible_gpu_count,
)
from utils.hardware import hardware as _hw_module
from core.inference.audio_codecs import AudioCodecManager
from io import StringIO
import structlog
@ -269,12 +253,7 @@ class InferenceBackend:
"""
Load any model: base, LoRA adapter, text, or vision.
"""
# max_seq_length=0 means "model default" for the GGUF/llama.cpp path,
# but Unsloth's FastLanguageModel.from_pretrained treats 0 literally --
# setting the model's context to 0 tokens, which triggers an assertion
# crash during generation (especially on ROCm/HIP where the async
# assert kernel raises a hardware exception instead of a Python error).
# Fall back to 2048 for the Unsloth/transformers path.
# GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it.
if max_seq_length <= 0:
max_seq_length = 2048
@ -311,12 +290,6 @@ class InferenceBackend:
}
# ── Audio model loading path ──────────────────────────
if (config.is_audio or config.is_vision) and _IS_ROCM_ENV:
raise RuntimeError(
f"Audio and vision model inference via Unsloth is not "
f"yet supported on AMD ROCm. Use GGUF inference instead."
)
if config.is_audio:
audio_type = config.audio_type
adapter_info = " (LoRA adapter)" if config.is_lora else ""
@ -547,84 +520,18 @@ class InferenceBackend:
else:
# Text model (or text LoRA adapter)
if _hw_module.IS_ROCM:
# On AMD ROCm two issues prevent the normal Unsloth path:
# 1. Unsloth's patched kernels (RoPE, attention) crash on
# HIP (_assert_async_cuda_kernel -> HSA_STATUS_ERROR).
# 2. bitsandbytes 4-bit matmul kernels trigger the same
# HIP assertion on MI300X (CDNA3 / gfx942).
# Fall back to plain transformers + PEFT in 16-bit, which
# works reliably. AMD GPUs typically have large VRAM so
# 16-bit is practical; GGUF inference remains the
# recommended path for memory-constrained setups.
logger.info(
"ROCm detected -- loading in 16-bit with plain "
"transformers (bitsandbytes 4-bit and Unsloth kernels "
"are not yet compatible with HIP)"
)
from transformers import AutoModelForCausalLM, AutoTokenizer
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = config.path, # Can be base model OR LoRA adapter path
max_seq_length = max_seq_length,
dtype = dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
token = hf_token if hf_token and hf_token.strip() else None,
trust_remote_code = trust_remote_code,
)
_load_kwargs = dict(
dtype = dtype or torch.bfloat16,
device_map = device_map,
token = hf_token if hf_token and hf_token.strip() else None,
trust_remote_code = trust_remote_code,
)
# Skip 4-bit on ROCm: bnb matmul kernels crash on HIP.
# Also resolve pre-quantized Unsloth model names (e.g.
# "unsloth/xxx-bnb-4bit") to their FP16 originals since
# loading a pre-quantized repo still triggers bnb codepaths.
def _resolve_fp16_base(name: str) -> str:
if not name:
return name
# Strip Unsloth quantization suffixes to get the FP16 model:
# "unsloth/Foo-unsloth-bnb-4bit" -> "unsloth/Foo"
# "unsloth/Foo-bnb-4bit" -> "unsloth/Foo"
# Order matters: try longer suffix first.
for suffix in ("-unsloth-bnb-4bit", "-bnb-4bit"):
if name.lower().endswith(suffix):
resolved = name[: -len(suffix)]
logger.info(
"Resolved pre-quantized base '%s' -> '%s' for ROCm 16-bit inference",
name,
resolved,
)
return resolved
return name
if config.is_lora and config.base_model:
# Load base model then apply adapter
_base = _resolve_fp16_base(config.base_model)
model = AutoModelForCausalLM.from_pretrained(
_base,
**_load_kwargs,
)
from peft import PeftModel
model = PeftModel.from_pretrained(model, config.path)
tokenizer = AutoTokenizer.from_pretrained(config.path)
else:
_path = _resolve_fp16_base(config.path)
model = AutoModelForCausalLM.from_pretrained(
_path,
**_load_kwargs,
)
tokenizer = AutoTokenizer.from_pretrained(config.path)
model.eval()
else:
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = config.path, # Can be base model OR LoRA adapter path
max_seq_length = max_seq_length,
dtype = dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
token = hf_token if hf_token and hf_token.strip() else None,
trust_remote_code = trust_remote_code,
)
# Apply inference optimization
FastLanguageModel.for_inference(model)
# Apply inference optimization
FastLanguageModel.for_inference(model)
self.models[model_name]["model"] = model
self.models[model_name]["tokenizer"] = tokenizer
@ -1047,13 +954,10 @@ class InferenceBackend:
)
# This modifies the tokenizer with the correct template
if get_chat_template is not None:
tokenizer = get_chat_template(
tokenizer,
chat_template = template_name,
)
else:
logger.info("Skipping Unsloth chat template (ROCm fallback)")
tokenizer = get_chat_template(
tokenizer,
chat_template = template_name,
)
else:
logger.info(
f"No registered Unsloth template for {self.active_model_name}, using tokenizer default"

101
studio/install_python_stack.py
View file

@ -43,6 +43,32 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = {
}
_PYTORCH_WHL_BASE = "https://download.pytorch.org/whl"
# bitsandbytes continuous-release_main wheels with the ROCm 4-bit GEMV fix
# (bnb PR #1887, post-0.49.2). bnb <= 0.49.2 NaNs at decode shape on every
# AMD GPU. Drop the pin once bnb 0.50+ ships on PyPI.
_BNB_ROCM_PRERELEASE_URLS: dict[str, str] = {
"x86_64": (
"https://github.com/bitsandbytes-foundation/bitsandbytes/releases/"
"download/continuous-release_main/"
"bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl"
),
"aarch64": (
"https://github.com/bitsandbytes-foundation/bitsandbytes/releases/"
"download/continuous-release_main/"
"bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl"
),
}
_BNB_ROCM_PYPI_FALLBACK = "bitsandbytes>=0.49.1"
def _bnb_rocm_prerelease_url() -> str | None:
"""Return the continuous-release_main bnb wheel URL for the current
architecture, or None when no pre-release wheel is available.
"""
arch = platform.machine().lower()
arch = {"amd64": "x86_64", "arm64": "aarch64"}.get(arch, arch)
return _BNB_ROCM_PRERELEASE_URLS.get(arch)
def _detect_rocm_version() -> tuple[int, int] | None:
"""Return (major, minor) of the installed ROCm stack, or None."""
@ -284,21 +310,37 @@ def _ensure_rocm_torch() -> None:
)
rocm_torch_ready = True
# Install bitsandbytes only when the venv has a ROCm-compatible torch
# (either already present or just installed). Avoids leaving an AMD
# bitsandbytes on top of a CPU/CUDA torch on hosts where the ROCm
# runtime is older than any published torch wheel. Uses
# --force-reinstall so an existing CPU/CUDA bitsandbytes is replaced
# by the AMD build during upgrades.
# Install bitsandbytes only when torch links against ROCm. Prefers the
# continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix) and falls
# back to PyPI when the pre-release URL is unreachable.
if rocm_torch_ready:
pip_install(
"bitsandbytes (AMD)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
"bitsandbytes>=0.49.1",
constrain = False,
)
_bnb_url = _bnb_rocm_prerelease_url()
_bnb_installed = False
if _bnb_url is not None:
_bnb_installed = pip_install_try(
"bitsandbytes (AMD, pre-release main)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
_bnb_url,
constrain = False,
)
if not _bnb_installed:
print(
_red(
" bnb pre-release unreachable; falling back to PyPI "
"(4-bit decode will be broken on ROCm)"
)
)
if not _bnb_installed:
pip_install(
"bitsandbytes (AMD)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
_BNB_ROCM_PYPI_FALLBACK,
constrain = False,
)
def _infer_no_torch() -> bool:
@ -593,6 +635,37 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]:
return cmd
def pip_install_try(
label: str,
*args: str,
constrain: bool = True,
) -> bool:
"""Like pip_install but returns False on failure instead of exiting.
For optional installs with a follow-up fallback.
"""
constraint_args: list[str] = []
if constrain and CONSTRAINTS.is_file():
constraint_args = ["-c", str(CONSTRAINTS)]
if USE_UV:
cmd = _build_uv_cmd(args) + constraint_args
else:
cmd = _build_pip_cmd(args) + constraint_args
if VERBOSE:
_step(_LABEL, f"{label}...", _dim)
result = subprocess.run(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
)
if result.returncode == 0:
return True
if VERBOSE and result.stdout:
print(result.stdout.decode(errors = "replace"))
return False
def pip_install(
label: str,
*args: str,

87
unsloth/models/rl.py
View file

@ -22,6 +22,8 @@ from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
import inspect
import os
import re
import sys
from contextlib import contextmanager
from unsloth_zoo.compiler import create_new_function
from unsloth_zoo.log import logger
from unsloth_zoo.logging_utils import PatchRLStatistics
@ -1947,6 +1949,83 @@ def patch_trl_rl_trainers():
return
def patch_trl_disable_gradient_checkpointing():
# TRL 1.0.0+ wraps generation in:
# with torch.no_grad(), disable_gradient_checkpointing(self.model, ...):
# The toggle exists only to suppress a cosmetic PyTorch warning
# ("None of the inputs have requires_grad=True"). Inside torch.no_grad()
# the gradient checkpointing state has no functional effect on the
# forward pass.
#
# On exit, the context manager calls model.gradient_checkpointing_enable()
# which dispatches to HuggingFace's generic implementation and overwrites
# Unsloth's custom `use_gradient_checkpointing="unsloth"` wrapper. For
# Gemma-4 (and likely other models) this corrupts the forward numerics
# enough to make GRPO KL divergence explode to ~10^12 at step 1.
#
# Replacing the context manager with a no-op preserves Unsloth's custom
# gradient checkpointing wrapper across generation/inference passes.
#
# Backwards compatibility:
# - trl < 1.0.0 (no disable_gradient_checkpointing): early return.
# - trl >= 1.0.0: noop is functionally equivalent for forward
# correctness. The only loss is a cosmetic warning being emitted
# by PyTorch when use_reentrant=True (which is exactly the warning
# TRL added the toggle to suppress in the first place).
try:
import trl.models.utils as _tmu
except ImportError:
return
if not hasattr(_tmu, "disable_gradient_checkpointing"):
return
if getattr(
_tmu.disable_gradient_checkpointing,
"_unsloth_noop_patched",
False,
):
return
@contextmanager
def _noop_disable_gradient_checkpointing(model, gradient_checkpointing_kwargs = None):
yield
_noop_disable_gradient_checkpointing._unsloth_noop_patched = True
_tmu.disable_gradient_checkpointing = _noop_disable_gradient_checkpointing
# Also rebind any trl.* module that already imported the symbol by
# reference, so the noop applies even when the trainer module cached the
# original at import time. We walk sys.modules dynamically rather than
# hardcoding a list, so this picks up every trainer that does
# `from ...models.utils import disable_gradient_checkpointing`
# (grpo, dpo, rloo, dppo, gfpo, grpo_with_replay_buffer, and any future
# TRL trainer module).
for _mod_name, _mod in list(sys.modules.items()):
if _mod is None or not _mod_name.startswith("trl."):
continue
try:
_bound = getattr(_mod, "disable_gradient_checkpointing", None)
except (AttributeError, ImportError):
continue
if _bound is None:
continue
try:
setattr(
_mod,
"disable_gradient_checkpointing",
_noop_disable_gradient_checkpointing,
)
except (AttributeError, TypeError):
pass
logger.warning_once(
"Unsloth: Patched trl.models.utils.disable_gradient_checkpointing with "
"a no-op to preserve Unsloth gradient checkpointing across TRL "
"generation passes."
)
return
def patch_trl_openenv():
for function in RL_ADDITIONAL_FUNCTIONS["openenv"]:
logger.info(f"Unsloth: Patching trl openenv with function: {function.__name__}")
@ -1981,6 +2060,14 @@ def patch_trl_vllm_generation():
def PatchFastRL(algorithm = None, FastLanguageModel = None):
if FastLanguageModel is not None:
PatchRL(FastLanguageModel)
# Install the disable_gradient_checkpointing noop BEFORE
# patch_trl_rl_trainers. patch_trl_rl_trainers imports extra trl.* trainer
# submodules while generating the compiled cache; any new trl.* modules
# imported after the sys.modules walk would keep their original (broken)
# binding of disable_gradient_checkpointing. Running the noop install
# first ensures the canonical trl.models.utils symbol is already replaced
# before those submodules bind it.
patch_trl_disable_gradient_checkpointing()
patch_trl_rl_trainers()
patch_trl_openenv()
patch_trl_vllm_generation()

35
unsloth/models/rl_replacements.py
View file

@ -855,9 +855,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
image_sizes_chunks = chunk_optional(image_sizes, B)
temperature = self.temperature
logit_softcapping = getattr(model.config, "final_logit_softcapping", 0)
if logit_softcapping is None:
logit_softcapping = 0
logit_softcapping = _unsloth_get_final_logit_softcapping(model.config)
logit_scale_multiply = getattr(model.config, "logit_scale", 0)
if logit_scale_multiply is None:
logit_scale_multiply = 0
@ -1004,11 +1002,38 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps_and_entropies)
def _unsloth_get_final_logit_softcapping(config):
"""Return final_logit_softcapping for a model config, falling back to the
nested text sub-config for composite models. Handles both:
- Gemma-4-style configs where the attribute lives on ``config.text_config``
- T5Gemma-style composite configs where the text sub-config is only
reachable via ``config.get_text_config()``
Returns 0 if unset, matching the previous behaviour.
"""
softcap = getattr(config, "final_logit_softcapping", None)
if softcap is None:
text_cfg = getattr(config, "text_config", None)
if text_cfg is None:
get_text_config = getattr(config, "get_text_config", None)
if callable(get_text_config):
try:
text_cfg = get_text_config()
except (TypeError, ValueError):
text_cfg = None
if text_cfg is not None and text_cfg is not config:
softcap = getattr(text_cfg, "final_logit_softcapping", None)
return 0 if softcap is None else softcap
grpo_compute_loss = RL_REPLACEMENTS["grpo_compute_loss"]
grpo_compute_loss_slow = RL_REPLACEMENTS["grpo_compute_loss_slow"]
UnslothEfficientGRPO = RL_REPLACEMENTS["UnslothEfficientGRPO"]
grpo_accumulated_loss = RL_REPLACEMENTS["grpo_accumulated_loss"]
grpo_update_SamplingParams = RL_REPLACEMENTS["grpo_update_SamplingParams"]
RL_PRE_ITEMS["grpo_trainer"].append(
inspect.getsource(_unsloth_get_final_logit_softcapping)
)
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_compute_loss))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(UnslothEfficientGRPO))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_accumulated_loss))
@ -1107,9 +1132,7 @@ def grpo_trainer_compute_loss(function_name, function):
input_ids = input_ids[:, -logits_to_keep:]
# Get logit softcapping and logit scale
logit_softcapping = getattr(model.config, "final_logit_softcapping", 0) # Gemma
if logit_softcapping is None:
logit_softcapping = 0
logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) # Gemma
logit_scale_multiply = getattr(model.config, "logit_scale", 0) # Cohere
if logit_scale_multiply is None:
logit_scale_multiply = 0