diff --git a/pyproject.toml b/pyproject.toml
index cad8580449..8eb6ea09a8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -64,6 +64,16 @@ cu121onlytorch211 = [
     "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10'",
     "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11'",
 ]
+cu118onlytorch212 = [
+    "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9'",
+    "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10'",
+    "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11'",
+]
+cu121onlytorch212 = [
+    "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23.post1-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9'",
+    "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23.post1-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10'",
+    "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23.post1-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11'",
+]
 cu118onlytorch220 = [
     "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.24%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9'",
     "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.24%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10'",
diff --git a/unsloth/models/dpo.py b/unsloth/models/dpo.py
index 3ae4d636f4..b7c7305bb3 100644
--- a/unsloth/models/dpo.py
+++ b/unsloth/models/dpo.py
@@ -12,11 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from transformers.utils.notebook import (
-    IntervalStrategy,
-    NotebookTrainingTracker,
-    NotebookProgressCallback,
-)
+try:
+    from transformers.utils.notebook import (
+        IntervalStrategy,
+        NotebookTrainingTracker,
+        NotebookProgressCallback,
+    )
+    HAS_NOTEBOOK = True
+except:
+    HAS_NOTEBOOK = False
+pass
 
 DPOTrainer_metrics = [
     "rewards/chosen",
@@ -101,13 +106,15 @@ pass
 
 
 def PatchDPOTrainer():
-    from transformers.trainer import is_in_notebook
-    if is_in_notebook():
-        # Patch DPO notebook printing
-        NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line
-        from transformers.trainer import DEFAULT_PROGRESS_CALLBACK
-        DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin
-        DEFAULT_PROGRESS_CALLBACK.on_log         = NotebookProgressCallback_on_log
+    if HAS_NOTEBOOK:
+        from transformers.trainer import is_in_notebook
+        if is_in_notebook():
+            # Patch DPO notebook printing
+            NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line
+            from transformers.trainer import DEFAULT_PROGRESS_CALLBACK
+            DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin
+            DEFAULT_PROGRESS_CALLBACK.on_log         = NotebookProgressCallback_on_log
+        pass
     pass
 pass
 
diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py
index 4aa634a4bd..d6c06edc9f 100644
--- a/unsloth/models/gemma.py
+++ b/unsloth/models/gemma.py
@@ -70,7 +70,7 @@ def GemmaDecoderLayer_fast_forward(
     padding_mask:         Optional[torch.LongTensor] = None,
     *args, **kwargs,
 ):
-    if False:#past_key_value is not None:
+    if past_key_value is not None:
         do_prefill = not hasattr(self.self_attn, "paged_attention")
 
         # Self Attention
@@ -267,6 +267,9 @@ class FastGemmaModel(FastLlamaModel):
         # Patch RMS Layernorm
         for name, module in model.named_modules():
             if isinstance(module, GemmaRMSNorm):
+                # Must be in float32
+                # https://github.com/keras-team/keras-nlp/blob/v0.8.2/keras_nlp/models/gemma/rms_normalization.py#L36
+                module = module.to(torch.float32)
                 module.weight += 1.0 # return output * (1 + self.weight)
                 if not hasattr(module, "variance_epsilon"):
                     module.variance_epsilon = module.eps # Gemma doesn't use variance_epsilon