diff --git a/pyproject.toml b/pyproject.toml index b3ee2d7aea..8ef5fcd240 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ triton = [ ] huggingface = [ - "unsloth_zoo>=2025.7.8", + "unsloth_zoo>=2025.7.9", "packaging", "tyro", "transformers>=4.51.3,!=4.47.0,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0", @@ -381,7 +381,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3", ] colab-new = [ - "unsloth_zoo>=2025.7.8", + "unsloth_zoo>=2025.7.9", "packaging", "tyro", "transformers>=4.51.3,!=4.47.0,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0", diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index df331fcd91..018e89c908 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -107,7 +107,7 @@ _cross_entropy_forward = triton.heuristics( def _chunked_cross_entropy_forward( logits_ptr , - logits_row_stride , + logits_row_stride : tl.constexpr, loss_ptr , logsumexp_ptr , labels_ptr , @@ -191,9 +191,9 @@ _chunked_cross_entropy_forward = triton.heuristics( def _cross_entropy_backward( logits_ptr , - logits_row_stride , + logits_row_stride : tl.constexpr, dloss_ptr , - dloss_row_stride , + dloss_row_stride : tl.constexpr, logsumexp_ptr , labels_ptr , VOCAB_SIZE : tl.constexpr, diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index fba7e56a84..ec45c6033b 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -19,9 +19,9 @@ from .utils import calculate_settings, torch_gpu_device @triton.jit def _rms_layernorm_forward( - Y, Y_row_stride, - X, X_row_stride, - W, W_row_stride, + Y, Y_row_stride : tl.constexpr, + X, X_row_stride : tl.constexpr, + W, W_row_stride : tl.constexpr, r, r_row_stride : tl.constexpr, n_cols : tl.constexpr, eps : tl.constexpr, @@ -54,10 +54,10 @@ pass def _rms_layernorm_backward( - dY, dY_row_stride, - dX, dX_row_stride, - X, X_row_stride, - W, W_row_stride, + dY, dY_row_stride : tl.constexpr, + dX, dX_row_stride : tl.constexpr, + X, X_row_stride : tl.constexpr, + W, W_row_stride : tl.constexpr, r, r_row_stride : tl.constexpr, # dW, dW_row_stride, n_cols : tl.constexpr, @@ -106,9 +106,9 @@ _rms_layernorm_backward = triton.heuristics( @triton.jit def _gemma_rms_layernorm_forward( - Y, Y_row_stride, - X, X_row_stride, - W, W_row_stride, + Y, Y_row_stride : tl.constexpr, + X, X_row_stride : tl.constexpr, + W, W_row_stride : tl.constexpr, r, r_row_stride : tl.constexpr, n_cols : tl.constexpr, eps : tl.constexpr, diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py index 1c981b3fc9..a06d81a512 100644 --- a/unsloth/kernels/rope_embedding.py +++ b/unsloth/kernels/rope_embedding.py @@ -19,9 +19,9 @@ from .utils import calculate_settings, torch_gpu_device ROPE_GROUP_SIZE : int = 4 def _rope_embedding( - Q, Q_row_stride, - cos, cos_row_stride, - sin, sin_row_stride, + Q, Q_row_stride: tl.constexpr, + cos, cos_row_stride: tl.constexpr, + sin, sin_row_stride: tl.constexpr, seqlen, head_dim : tl.constexpr, n_heads : tl.constexpr, diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 645319d423..1470068e76 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -124,7 +124,7 @@ if DEVICE_TYPE == "xpu": (index := torch.xpu.device(i).idx) : ctypes.c_void_p(torch._C._xpu_getCurrentRawStream(index)) for i in range(DEVICE_COUNT) } - XPU_STREAMS = [None] * (max(_XPU_STREAMS.keys()) + 1) + XPU_STREAMS = [None] * (max(_XPU_STREAMS.keys()) + 1) WEIGHT_BUFFERS = [None] * (max(_XPU_STREAMS.keys()) + 1) ABSMAX_BUFFERS = [None] * (max(_XPU_STREAMS.keys()) + 1) for k, v in _XPU_STREAMS.items(): @@ -143,7 +143,7 @@ else: for k, v in _CUDA_STREAMS.items(): CUDA_STREAMS[k] = v CUDA_STREAMS = tuple(CUDA_STREAMS) del _CUDA_STREAMS - +pass # Bitsandbytes operations ctypes_c_int = ctypes.c_int @@ -172,12 +172,16 @@ else: cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4 cgemm_4bit_inference_naive_fp16 = bnb.functional.lib.cgemm_4bit_inference_naive_fp16 cgemm_4bit_inference_naive_bf16 = bnb.functional.lib.cgemm_4bit_inference_naive_bf16 +pass torch_mm = torch.mm torch_mv = torch.mv torch_matmul = torch.matmul torch_addmm = torch.addmm torch_empty = torch.empty +torch_float32 = torch.float32 +torch_float16 = torch.float16 +torch_bfloat16 = torch.bfloat16 def QUANT_STATE(W): return getattr(W, "quant_state", None) @@ -283,7 +287,7 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM: else: assert(out.shape == shape) assert(out.dtype == dtype) - out_absmax = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + out_absmax = torch_empty(n_elements_absmax, dtype = torch_float32, device = device, requires_grad = False) pass # NF4 dequantization of statistics @@ -296,7 +300,7 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM: out_absmax += offset # Dequantize W - fx = cdequantize_blockwise_fp16_nf4 if dtype == torch.float16 else \ + fx = cdequantize_blockwise_fp16_nf4 if dtype == torch_float16 else \ cdequantize_blockwise_bf16_nf4 fx(get_ptr(None), get_ptr(W), ptr_out_absmax, get_ptr(out), ctypes_c_int(blocksize), ctypes_c_int(out.numel()), XPU_STREAM,) @@ -346,7 +350,7 @@ elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: ABSMAX_BUFFER = ABSMAX_BUFFERS[device_index] if WEIGHT_BUFFER is None: WEIGHT_BUFFERS[device_index] = WEIGHT_BUFFER = torch_empty(size, dtype = dtype, device = device, requires_grad = False) - ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch_empty(n_elements_absmax, dtype = torch_float32, device = device, requires_grad = False) if size > WEIGHT_BUFFER.numel(): WEIGHT_BUFFER.resize_(size) if n_elements_absmax > ABSMAX_BUFFER.numel(): ABSMAX_BUFFER.resize_(n_elements_absmax) @@ -359,7 +363,7 @@ elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: else: assert(out.shape == shape) assert(out.dtype == dtype) - out_absmax = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + out_absmax = torch_empty(n_elements_absmax, dtype = torch_float32, device = device, requires_grad = False) pass # NF4 dequantization of statistics @@ -372,7 +376,7 @@ elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: out_absmax += offset # Dequantize W - fx = cdequantize_blockwise_fp16_nf4 if dtype == torch.float16 else \ + fx = cdequantize_blockwise_fp16_nf4 if dtype == torch_float16 else \ cdequantize_blockwise_bf16_nf4 fx(get_ptr(None), get_ptr(W), ptr_out_absmax, get_ptr(out), ctypes_c_int(blocksize), ctypes_c_int(out.numel()), CUDA_STREAM,) @@ -413,7 +417,7 @@ else: else: assert(out.shape == shape) assert(out.dtype == dtype) - out_absmax = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + out_absmax = torch_empty(n_elements_absmax, dtype = torch_float32, device = device, requires_grad = False) # Do dequantization ptr_out_absmax = get_ptr(out_absmax) @@ -423,7 +427,7 @@ else: ) out_absmax += offset - fx = cdequantize_blockwise_fp16_nf4 if dtype == torch.float16 else \ + fx = cdequantize_blockwise_fp16_nf4 if dtype == torch_float16 else \ cdequantize_blockwise_bf16_nf4 fx(get_ptr(None), get_ptr(W), ptr_out_absmax, get_ptr(out), ctypes_c_int(blocksize), ctypes_c_int(out.numel()),) @@ -488,7 +492,7 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM: ldb = ctypes_c_int32(ldb) ldc = ctypes_c_int32(ldc) - df = torch_empty(absmax.shape, dtype = torch.float32, device = device) + df = torch_empty(absmax.shape, dtype = torch_float32, device = device) with torch_gpu_device(device): cdequantize_blockwise_fp32( get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), @@ -497,7 +501,7 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM: df += offset absmax = df - fx = cgemm_4bit_inference_naive_fp16 if dtype == torch.float16 else \ + fx = cgemm_4bit_inference_naive_fp16 if dtype == torch_float16 else \ cgemm_4bit_inference_naive_bf16 blocksize = ctypes_c_int32(blocksize) @@ -559,7 +563,7 @@ elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: ldb = ctypes_c_int32(ldb) ldc = ctypes_c_int32(ldc) - df = torch_empty(absmax.shape, dtype = torch.float32, device = device) + df = torch_empty(absmax.shape, dtype = torch_float32, device = device) with torch_gpu_device(device): cdequantize_blockwise_fp32( get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), @@ -568,8 +572,8 @@ elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: df += offset absmax = df - fx = cgemm_4bit_inference_naive_fp16 if dtype == torch.float16 else \ - cgemm_4bit_inference_naive_bf16 + fx = cgemm_4bit_inference_naive_fp16 if dtype == torch_float16 else \ + cgemm_4bit_inference_naive_bf16 blocksize = ctypes_c_int32(blocksize) fx(m, n, k, get_ptr(X), get_ptr(W), get_ptr(absmax), get_ptr(stats), get_ptr(out), @@ -580,7 +584,7 @@ elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: pass else: def fast_gemv(X, W, quant_state, out = None): - if quant_state is None: return torch.matmul(X, W, out = out) + if quant_state is None: return torch_matmul(X, W, out = out) # For fast X @ W where seq_len == 1 # From https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L1469 _, q_len, hd = X.shape @@ -626,7 +630,7 @@ else: ldb = ctypes_c_int32(ldb) ldc = ctypes_c_int32(ldc) - df = torch_empty(absmax.shape, dtype = torch.float32, device = device) + df = torch_empty(absmax.shape, dtype = torch_float32, device = device) cdequantize_blockwise_fp32( get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), ctypes_c_int(blocksize2), ctypes_c_int(df.numel()), @@ -634,8 +638,8 @@ else: df += offset absmax = df - fx = cgemm_4bit_inference_naive_fp16 if dtype == torch.float16 else \ - cgemm_4bit_inference_naive_bf16 + fx = cgemm_4bit_inference_naive_fp16 if dtype == torch_float16 else \ + cgemm_4bit_inference_naive_bf16 blocksize = ctypes_c_int32(blocksize) fx(m, n, k, get_ptr(X), get_ptr(W), get_ptr(absmax), get_ptr(stats), get_ptr(out), diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index c445e98b07..bce6bb616d 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2025.7.7" +__version__ = "2025.7.8" __all__ = [ "SUPPORTS_BFLOAT16", @@ -431,7 +431,7 @@ if DEVICE_TYPE == "cuda": "Unsloth: If you want to finetune Gemma 2, upgrade flash-attn to version 2.6.3 or higher!\n"\ "Newer versions support faster and less memory usage kernels for Gemma 2's attention softcapping!\n"\ "To update flash-attn, do the below:\n"\ - '\npip install --no-deps --upgrade "flash-attn>=2.6.3"' + '\npip install --no-deps --no-build-isolation --upgrade "flash-attn>=2.6.3"' ) except: print( diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 5597995b05..0f1924579a 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -59,18 +59,6 @@ pass if HAS_FLASH_ATTENTION_SOFTCAPPING: from flash_attn import flash_attn_func -# [TODO] We must randomnly use torch.compile? -# Gemma 2 uses double RMS Layernorms, so the backward passes should not overwrite the gradients! -@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) -def fast_rms_layernorm_gemma2_compiled(layernorm, X, gemma = True): - old_dtype = X.dtype - X = X.float() - X = X * torch.rsqrt(X.square().mean(-1, keepdim = True) + layernorm.eps) * \ - (1.0 + layernorm.weight.float()) - return X.to(old_dtype) -pass - - # Logit softcapping def Gemma2Attention_fast_forward( self, diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 5bbf4c76a6..2436db4ff4 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -85,6 +85,12 @@ from unsloth_zoo.vllm_utils import ( return_lora_modules, ) +try: + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_compiler_set_stance = None +pass + def unsloth_base_fast_generate( self, *args, @@ -756,7 +762,8 @@ class FastBaseModel: # Must enable returning logits os.environ["UNSLOTH_RETURN_LOGITS"] = "1" # Turn off skip guards and set stance to default - torch.compiler.set_stance(stance = "default", skip_guard_eval_unsafe = False) + if torch_compiler_set_stance is not None: + torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) return model pass @@ -804,7 +811,8 @@ class FastBaseModel: # Can re-enable not returning logits os.environ["UNSLOTH_RETURN_LOGITS"] = "0" # Turn off skip guards and set stance to default - torch.compiler.set_stance(stance = "default", skip_guard_eval_unsafe = False) + if torch_compiler_set_stance is not None: + torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) return model pass pass