Repository: huggingface/transformers
Issue category: Feature Request

Add support for models trained using MLX
### Feature request

It would be great if model weights trained and saved through MLX could be easily loaded using the Transformers library. This way, MLX users and Transformers users can more freely collaborate when making open source models.


### Motivation

Currently, Transformers only supports safetensors files with metadata for format in `["pt", "tf", "flax"]`. If Transformers can add `"mlx"` to the list, and if MLX can add `metadata={"format":"mlx"}` when saving safetensors, this can be achieved.

Here is a related issue I created on the `mlx_lm` side:
https://github.com/ml-explore/mlx/issues/743#issuecomment-1965427589

Here is a related PR on the `mlx_lm` repo: https://github.com/ml-explore/mlx-examples/pull/496

### Your contribution

I can send a PR to modify the exact code shortly after.


Functions to edit:
src/transformers/modeling_utils.py:PreTrainedModel.from_pretrained

Patch:
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 505c9cb45950cb..ca98c64a29a823 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -3297,9 +3297,12 @@ def from_pretrained(
             elif metadata.get("format") == "flax":
                 from_flax = True
                 logger.info("A Flax safetensors file is being loaded in a PyTorch model.")
+            elif metadata.get("format") == "mlx":
+                # This is a mlx file, we assume weights are compatible with pt
+                pass
             else:
                 raise ValueError(
-                    f"Incompatible safetensors file. File metadata is not ['pt', 'tf', 'flax'] but {metadata.get('format')}"
+                    f"Incompatible safetensors file. File metadata is not ['pt', 'tf', 'flax', 'mlx'] but {metadata.get('format')}"
                 )
 
         from_pt = not (from_tf | from_flax)

Test patch: diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
index d0db5031e8b7a0..acd60fbbea3119 100755
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -1256,6 +1256,26 @@ def test_modifying_model_config_causes_warning_saving_generation_config(self):
             self.assertEqual(len(logs.output), 1)
             self.assertIn("Your generation config was originally created from the model config", logs.output[0])
 
+    @require_safetensors
+    def test_model_from_pretrained_from_mlx(self):
+        from safetensors import safe_open
+
+        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-mistral-mlx")
+        self.assertIsNotNone(model)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir, safe_serialization=True)
+            with safe_open(os.path.join(tmp_dir, "model.safetensors"), framework="pt") as f:
+                metadata = f.metadata()
+                self.assertEqual(metadata.get("format"), "pt")
+            new_model = AutoModelForCausalLM.from_pretrained(tmp_dir)
+
+        input_ids = torch.randint(100, 1000, (1, 10))
+        with torch.no_grad():
+            outputs = model(input_ids)
+            outputs_from_saved = new_model(input_ids)
+            self.assertTrue(torch.allclose(outputs_from_saved["logits"], outputs["logits"]))
+
 
 @slow
 @require_torch