gemma

unsloth/kernels/cross_entropy_loss.py

@@ -168,7 +168,7 @@ def fast_cross_entropy_loss(logits, labels):
             "Unsloth's Github page if you want a faster and more memory efficient kernel!"
         )
         loss = slow_cross_entropy_loss(
-            logits.view(batch*seq_len, d), # Must cast to float32 for numerical stability
+            logits.float().view(batch*seq_len, d), # Must cast to float32 for numerical stability
             labels.view(-1),
         )
         return loss

unsloth/models/loader.py

@@ -14,7 +14,6 @@
 
 from .llama import FastLlamaModel, logger
 from .mistral import FastMistralModel
-from .gemma import FastGemmaModel
 from transformers import AutoConfig
 from transformers import __version__ as transformers_version
 from peft import PeftConfig, PeftModel
@@ -26,6 +25,8 @@ major, minor = transformers_version.split(".")[:2]
 major, minor = int(major), int(minor)
 SUPPORTS_FOURBIT = (major > 4) or (major == 4 and minor >= 37)
 SUPPORTS_GEMMA   = (major > 4) or (major == 4 and minor >= 38)
+if SUPPORTS_GEMMA:
+    from .gemma import FastGemmaModel
 del major, minor