updates

Author: Goekdeniz-Guelmez

.gitignore

@@ -25,3 +25,5 @@ var/
 # virtual environment
 .myenv
 .venv
+
+.vscode

mlx_vlm/lora.py

@@ -12,6 +12,7 @@
     find_all_linear_names,
     get_peft_model,
     unfreeze_modules,
+    supported_for_training
 )
 from .utils import load, load_image_processor
 
@@ -25,9 +26,8 @@ def main(args):
         args.model_path, processor_config={"trust_remote_code": True}
     )
 
-    unsupported_for_training = {"lfm2-vl", "", ""}
     model_type = getattr(getattr(model, "config", None), "model_type", None)
-    if model_type in unsupported_for_training:
+    if model_type not in supported_for_training:
         raise ValueError(
             f"{Colors.FAIL}Model type {model_type} not supported for training. "
             f"Please choose a different model or remove it from the unsupported list.{Colors.ENDC}"

mlx_vlm/models/qwen2_vl/qwen2_vl.py

@@ -154,6 +154,9 @@ def __call__(
         return logits
 
     def sanitize(self, weights):
+        weights = {
+            k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
+        }
         def transform_key(key):
             if "vision_tower" not in key:
                 key = key.replace("visual", "vision_tower")

mlx_vlm/trainer/__init__.py

@@ -8,4 +8,5 @@
     get_peft_model,
     print_trainable_parameters,
     Colors,
+    supported_for_training
 )

mlx_vlm/trainer/utils.py

@@ -20,6 +20,7 @@ class Colors:
     BOLD = '\033[1m'
     UNDERLINE = '\033[4m'
 
+supported_for_training = {"qwen2_vl"}
 
 def grad_checkpoint(layer):
     """