feat: add mlx model and trainer

[JINO-ROHIT]

hello everyone!

this PR aims to integrate mlx support in unsloth with minimal changes.

1. ive tried to keep the PR as compact as possible and make use of the existing mlx utilities.
2. ive also had to make some patches on the unsloth-zoo code files, should i raise a seperate PR for that?

im attaching below a sample alpaca training run script to get this working.

from unsloth.models.mlx_model import FastMLXModel
model, tokenizer = FastMLXModel.from_pretrained("mlx-community/Llama-3.2-3B-Instruct-4bit")


from datasets import load_dataset
dataset = load_dataset("mlabonne/FineTome-Alpaca-100k", split="train")

system_message = """You are an assistant."""
def create_conversation(sample):
  return {
    "messages": [
      {"role": "system", "content": system_message},
      {"role": "user", "content": sample["instruction"]}, # human
      {"role": "assistant", "content": sample["output"]} # model
    ]
  }

dataset = dataset.map(create_conversation, remove_columns=dataset.features,batched=False)
dataset = dataset.train_test_split(0.1)

from mlx_lm.tuner import datasets

configs = {
    "mask_prompt": False,
    "prompt_feature": "prompt",
    "text_feature": "text",
    "completion_feature": "completion",
    "chat_feature": "messages",
}

train_set = datasets.create_dataset(
    dataset["train"],
    tokenizer,
    configs
)

val_set = datasets.create_dataset(
    dataset["test"],
    tokenizer,
    configs
)


FastMLXModel.train(
    model,
    train_set,
    val_set,
    iterations = 2
)

[gemini-code-assist]

Summary of Changes

Hello @JINO-ROHIT, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly expands Unsloth's hardware compatibility by integrating MLX support, primarily benefiting users on Apple Silicon devices. The changes enable the loading and fine-tuning of models using the MLX framework, providing an optimized experience for Mac users. This integration is achieved through new dedicated MLX model and trainer classes, alongside conditional logic to ensure seamless operation across different hardware platforms.

Highlights

• MLX Integration: Introduced support for MLX (Machine Learning eXchange) models and training, specifically targeting Apple Silicon (MPS) devices.
• New MLX Model and Trainer: Added a new file unsloth/models/mlx_model.py which defines FastMLXModel for loading and training MLX-compatible models, along with MLXTrainer, MLXTrainingArguments, and MLXLoraConfig.
• Conditional Imports and Device Detection: Modified core Unsloth files (__init__.py, device_type.py, models/__init__.py) to conditionally import MLX-related modules and correctly detect and handle MPS devices, ensuring MLX functionality is enabled only when appropriate.
• Dependency Management: Updated pyproject.toml to include mlx and mlx-lm as dependencies for macOS (arm64) platforms, enabling MLX support for Apple Silicon users.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces MLX support to unsloth, enabling training on Apple Silicon. The changes are well-structured and mostly confined to conditional imports and a new mlx_model.py module. I've identified a few areas for improvement to enhance code clarity and reduce redundancy. Specifically, I've suggested removing a redundant conditional block in device_type.py, simplifying a dataclass-to-dictionary conversion, and removing a duplicate check for lora_config in the new MLX model file. Overall, this is a great addition to the library.

[gemini-code-assist]

This elif block is redundant. Since DEVICE_TYPE_TORCH is initialized with the value of DEVICE_TYPE, if DEVICE_TYPE_TORCH is 'mps', this block just assigns 'mps' back to it. This block can be safely removed to improve code clarity. Also, there is a trailing whitespace on line 70.

[gemini-code-assist]

The to_dict method can be simplified by using dataclasses.asdict. This will make the code more concise and less prone to errors if new fields are added to the dataclass, as you won't need to update this method manually.

To use it, you'll also need to update the import on line 4 to:

from dataclasses import dataclass, asdict

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

[gemini-code-assist]

This check for lora_config being None is redundant. The prepare_model_for_training method already handles the case where lora_config is None by creating a default MLXLoraConfig. Removing this duplication will make the code cleaner.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c623f4cd3d

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Base modules never imported on CUDA/HIP/XPU

The core imports are now guarded by elif DEVICE_TYPE != "mps":, but that condition is part of the same if/elif chain as the preceding CUDA/HIP/XPU branches. Because one of those earlier branches always matches on supported GPUs, this block never executes, so import unsloth no longer brings in FastLanguageModel, __version__, trainer/chat helpers, or the TRL patch on any CUDA/ROCm/Intel system—the chain exits before reaching these imports—leading to AttributeError/missing functionality for all existing users.

Useful? React with 👍 / 👎.

[shimmyshimmer]

Thank you so much for your PR we'll take a look asap! :)

[JustinWick]

Shouldn't that be M-Series Apple Silicon, Apple Macintosh Silicon or Apple Mac Silicon? I don't think it's ever branded "MAC".

[JINO-ROHIT]

hello @shimmyshimmer sorry did you get some time to review this? id love to iterate on feedbacks and push this PR forward

[shimmyshimmer]

hello @shimmyshimmer sorry did you get some time to review this? id love to iterate on feedbacks and push this PR forward

While the PR itself is fantastic (thank you!!), there arent any optimizations at the moment, were discussing whether we should proceed as is or, spend more time on optimizations

[b-straub]

Being only an interested consumer of Unsloth models I might be wrong, but to my knowledge the fine-tuning performance gains are also heavily related to the custom Triton kernels. Shouldn’t the optimization strategies behind those kernels be reimplemented for MLX when possible? (I understand Triton kernels can’t be directly ported since they compile to CUDA PTX, but the underlying approaches like fused attention and memory-efficient backward passes could potentially be implemented using MLX’s Metal primitives.)

[JINO-ROHIT]

hello @shimmyshimmer sorry did you get some time to review this? id love to iterate on feedbacks and push this PR forward

While the PR itself is fantastic (thank you!!), there arent any optimizations at the moment, were discussing whether we should proceed as is or, spend more time on optimizations

of course, ill wait to hear on further updates

[JINO-ROHIT]

Being only an interested consumer of Unsloth models I might be wrong, but to my knowledge the fine-tuning performance gains are also heavily related to the custom Triton kernels. Shouldn’t the optimization strategies behind those kernels be reimplemented for MLX when possible? (I understand Triton kernels can’t be directly ported since they compile to CUDA PTX, but the underlying approaches like fused attention and memory-efficient backward passes could potentially be implemented using MLX’s Metal primitives.)

i think mlx should already have its own set of pre-existing optimizations, but sure we could also possibly look into further improvements, although im not quite sure how difficult/easy itd be to make new metal kernels and integrate within.

[b-straub]

think mlx should already have its own set of pre-existing optimizations, but sure we could also possibly look into further improvements, although im not quite sure how difficult/easy itd be to make new metal kernels and integrate within.

At least should be easier now https://ml-explore.github.io/mlx/build/html/dev/custom_metal_kernels.html

[JINO-ROHIT]

sure, i meant the overall complexity of metal programming + wrtiting optimized kernels on top of it

[Manan17]

@JINO-ROHIT can you please upload the changes you made in unsloth_zoo in order to run this.

[JINO-ROHIT]

sure ill do it

[xpatronum]

@JINO-ROHIT any updates, plz ?

[braduck]

Who can Resolve the conflicts and merge?

[negative0]

Eagerly waiting for this 🥹

[danielhanchen]

Sorry we had to rebase the PR - apologies!

[step21]

Wouldn't it make more sense to base this just on gpu support in f.e. pytorch then on MLX? this would enable usage of the same models.

[RukshanJS]

what happened with this?