Multipack simplify for Mixtral (#1142)

2024-01-18 16:23:49 -05:00
parent 1d70f24b50
commit 6910e6a8ca
11 changed files with 201 additions and 430 deletions
--- a/tests/e2e/patched/test_model_patches.py
+++ b/tests/e2e/patched/test_model_patches.py
@@ -52,11 +52,7 @@ class TestModelPatches(unittest.TestCase):
        model, _ = load_model(cfg, tokenizer, inference=cli_args.inference)

        assert (
-            "axolotl.monkeypatch.mixtral.modeling_mixtral"
-            in model.model.layers[0].self_attn.__class__.__module__
-        )
-        assert (
-            "MixtralMultipackFlashAttention2"
+            "MixtralFlashAttention2"
            in model.model.layers[0].self_attn.__class__.__name__
        )