llama test

2025-01-12 17:30:19 -05:00
parent 888cd9407f
commit f99cae0e7b
2 changed files with 4 additions and 3 deletions
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -387,12 +387,13 @@ class ModelLoader:
            self.patch_attention()

        if self.cfg.model_config_type == "llama":
-            from axolotl.monkeypatch.trainer_grad_accum import (  # patch_flash_attention_forward,
+            from axolotl.monkeypatch.trainer_grad_accum import (
+                patch_flash_attention_forward,
                patch_forward_for_ga,
                patch_training_step_for_ga,
            )

-            # patch_flash_attention_forward()
+            patch_flash_attention_forward()
            patch_forward_for_ga()
            patch_training_step_for_ga()