diff --git a/install.sh b/install.sh index 2e347a8bc5..cc8721710e 100755 --- a/install.sh +++ b/install.sh @@ -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 diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index d7bc3ad4c7..4c140013a0 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -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" diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index af77d7a811..a046f8f892 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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, diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 5651a7da41..f444f6bd37 100755 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -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() diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 2544afe82e..93a7f89bcb 100755 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -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