Vision support

unsloth/__init__.py

@@ -55,7 +55,7 @@
 pass
 
 # Reduce VRAM usage by reducing fragmentation
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,roundup_power2_divisions:[64:128,256:64,>:32]"
 
 # Hugging Face Hub faster downloads
 if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ:

unsloth/kernels/__init__.py

@@ -42,6 +42,7 @@
     apply_lora_mlp_geglu_approx,
     apply_lora_qkv,
     apply_lora_o,
+    fast_lora_forward,
 )
 from .utils import fast_dequantize, fast_gemv, QUANT_STATE, fast_linear_forward, matmul_lora
 

unsloth/kernels/fast_lora.py

@@ -410,3 +410,81 @@ def apply_lora_o(self, X):
     O = LoRA_W.apply(X, OW, OW_quant, OA, OB, OS)
     return O
 pass
+
+
+IDENTITY_DROPOUT = torch.nn.Identity
+@torch._disable_dynamo
+def fast_lora_forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+    raise NotImplementedError(
+        "Unsloth: Currently not supported yet - reshaping done incorrectly"
+    )
+    self._check_forward_args(x, *args, **kwargs)
+    adapter_names = kwargs.pop("adapter_names", None)
+
+    if self.disable_adapters:
+        if self.merged:
+            self.unmerge()
+        result = self.base_layer(x, *args, **kwargs)
+    elif adapter_names is not None:
+        result = self._mixed_batch_forward(x, *args, adapter_names=adapter_names, **kwargs)
+    elif self.merged:
+        result = self.base_layer(x, *args, **kwargs)
+    else:
+        # Fastpath
+        if len(self.active_adapters) == 1:
+            active_adapter = self.active_adapters[0]
+            if active_adapter not in self.lora_A.keys(): return self.base_layer(x, *args, **kwargs)
+
+            dropout = self.lora_dropout[active_adapter]
+            if isinstance(dropout, IDENTITY_DROPOUT) and not self.use_dora[active_adapter]:
+                lora_A = self.lora_A[active_adapter].weight
+                lora_B = self.lora_B[active_adapter].weight
+                scaling = self.scaling[active_adapter]
+                W = self.base_layer.weight
+                return LoRA_W.apply(x, W, QUANT_STATE(W), lora_A, lora_B, scaling)
+            pass
+        pass
+
+        result = self.base_layer(x, *args, **kwargs)
+        # As per Tim Dettmers, for 4bit, we need to defensively clone here.
+        # The reason is that in some cases, an error can occur that backprop
+        # does not work on a manipulated view. This issue may be solved with
+        # newer PyTorch versions but this would need extensive testing to be
+        # sure.
+        result = result.clone()
+
+        for active_adapter in self.active_adapters:
+            if active_adapter not in self.lora_A.keys():
+                continue
+            lora_A = self.lora_A[active_adapter]
+            lora_B = self.lora_B[active_adapter]
+            dropout = self.lora_dropout[active_adapter]
+            scaling = self.scaling[active_adapter]
+
+            requires_conversion = not torch.is_autocast_enabled()
+            if requires_conversion:
+                expected_dtype = result.dtype
+                x = x.to(lora_A.weight.dtype)
+
+            if not self.use_dora[active_adapter]:
+                result = result + lora_B(lora_A(dropout(x))) * scaling
+            else:
+                if isinstance(dropout, torch.nn.Identity) or not self.training:
+                    base_result = result
+                else:
+                    x = dropout(x)
+                    base_result = None
+
+                result = result + self.lora_magnitude_vector[active_adapter](
+                    x,
+                    lora_A=lora_A,
+                    lora_B=lora_B,
+                    scaling=scaling,
+                    base_layer=self.get_base_layer(),
+                    base_result=base_result,
+                )
+            if requires_conversion:
+                result = result.to(expected_dtype)
+
+    return result
+pass

unsloth/kernels/rms_layernorm.py

@@ -57,6 +57,7 @@ def _rms_layernorm_forward(
 @triton.jit
 def _rms_layernorm_backward(
     dY, dY_row_stride,
+    dX, dX_row_stride,
     X,   X_row_stride,
     W,   W_row_stride,
     r,   r_row_stride,
@@ -78,6 +79,9 @@ def _rms_layernorm_backward(
     X  += row_idx *  X_row_stride
     r  += row_idx *  r_row_stride
 
+    if GEMMA: dX += row_idx * dY_row_stride
+    else:     dX = dY
+
     dY_row = tl.load(dY + col_offsets, mask = mask, other = 0).to(tl.float32)
     X_row  = tl.load(X  + col_offsets, mask = mask, other = 0).to(tl.float32)
     W_row  = tl.load(W  + col_offsets, mask = mask, other = 0).to(tl.float32)
@@ -91,7 +95,7 @@ def _rms_layernorm_backward(
 
     rowsum_dY_normed = tl.sum(dY_W * normed, axis = 0)
     output = inv_var/n_cols * (n_cols*dY_W - normed*rowsum_dY_normed)
-    tl.store(dY + col_offsets, output, mask = mask)
+    tl.store(dX + col_offsets, output, mask = mask)
 pass
 
 
@@ -172,9 +176,11 @@ def backward(ctx, dY : torch.Tensor):
         n_cols : int
         n_rows, n_cols = dY.shape
         # dW = X
+        dX = torch.empty_like(dY, device = "cuda:0") if ctx.GEMMA else dY
 
         _rms_layernorm_backward[(n_rows,)](
             dY, dY.stride(0),
+            dX, dX.stride(0),
             X,  X .stride(0),
             W,  W .stride(0),
             r,  r .stride(0),
@@ -184,7 +190,7 @@ def backward(ctx, dY : torch.Tensor):
             BLOCK_SIZE = ctx.BLOCK_SIZE,
             num_warps  = ctx.num_warps,
         )
-        dX = dY.view(*shape)
+        dX = dX.view(*shape)
         return dX, None, None, None
     pass
 pass

unsloth/models/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .loader  import FastLanguageModel
+from .loader  import FastLanguageModel, FastVisionModel
 from .llama   import FastLlamaModel
 from .mistral import FastMistralModel
 from .qwen2   import FastQwen2Model