adding multi token input and correct cache handling in ssm step

2025-09-18 01:50:19 +08:00 · 2024-10-22 20:44:23 +02:00
parent 5326d9373a
commit 758597eaa8
2 changed files with 251 additions and 137 deletions
--- a/llms/mlx_lm/tuner/utils.py
+++ b/llms/mlx_lm/tuner/utils.py
@@ -146,6 +146,8 @@ def linear_to_lora_layers(
    elif model.model_type == "mamba2":
        keys = set(
            [
+                "mixer.in_proj",
+                "mixer.out_proj",
            ]
        )
    else: