Skip to main content

Use MLX in Strands Agents

Project description

strands-mlx

PyPI version Interactive Agent

Running Strands Agents locally on Apple Silicon - inference, fine-tuning, vision in Python

MLX provider for Strands Agents with LoRA training pipelines.


Requirements: Python ≤3.13, macOS/Linux

# Text models
pip install strands-mlx

# With vision/audio/video
pip install "strands-mlx[vision]"

# Environment setup
uv venv --python 3.13 && source .venv/bin/activate

# Install dependencies
uv pip install strands-agents strands-mlx strands-agents-tools

Quick Start

agent.py

from strands import Agent
from strands_mlx import MLXModel
from strands_tools import calculator

model = MLXModel(model_id="mlx-community/Qwen3-1.7B-4bit")
agent = Agent(model=model, tools=[calculator])

agent("What is 29 * 42?")
# Run with uv
uv run agent.py

Architecture

graph LR
    A[Agent Conversations] -->|MLXSessionManager| B[Training Data JSONL]
    B -->|dataset_splitter| C[train/valid/test]
    C -->|mlx_trainer| D[LoRA Adapter]
    D -->|MLXModel| E[Domain Expert Agent]
    E -.->|Continuous Learning| A
    
    style A fill:#e1f5ff
    style E fill:#d4edda
    style D fill:#fff3cd

The complete training cycle: Agents collect their own training data → fine-tune themselves → become domain experts → continue learning.


Train Your Own Model

4 steps: Collect → Split → Train → Use

1. Collect Training Data

from strands import Agent
from strands_tools import calculator
from strands_mlx import MLXSessionManager, dataset_splitter, mlx_trainer

session = MLXSessionManager(session_id="my_training", storage_dir="./dataset")
agent = Agent(model=model, tools=[calculator, dataset_splitter, mlx_trainer], session_manager=session)

# Have conversations - auto-saved to JSONL
agent("Teach me about quantum computing")
agent("Calculate 15 * 7")

# Saved to: ./dataset/my_training.jsonl

2. Split Dataset

agent.tool.dataset_splitter(
    input_path="./dataset/my_training.jsonl"
)
# Creates train.jsonl, valid.jsonl, test.jsonl (80/10/10 split)

3. Train with LoRA

agent.tool.mlx_trainer(
    action="train",
    config={
        "model": "mlx-community/Qwen3-1.7B-4bit",
        "data": "./dataset/my_training",
        "adapter_path": "./adapter",
        "iters": 200,
        "learning_rate": 1e-5,
        "batch_size": 1
    }
)

4. Use Trained Model

from strands import Agent
from strands_mlx import MLXModel

trained = MLXModel("mlx-community/Qwen3-1.7B-4bit", adapter_path="./adapter")
agent = Agent(model=trained)

agent("Explain quantum computing")  # Uses trained knowledge!

Vision Models

from strands_mlx import MLXVisionModel

model = MLXVisionModel(model_id="mlx-community/Qwen2-VL-2B-Instruct-4bit")
agent = Agent(model=model)

agent("Describe: <image>photo.jpg</image>")
agent("Transcribe: <audio>speech.wav</audio>")
agent("What happens: <video>clip.mp4</video>")

Training Tools

Tool Purpose
mlx_trainer Background LoRA training
dataset_splitter Split JSONL → train/valid/test
validate_training_data Check format & token counts
mlx_invoke Runtime model switching
mlx_vision_invoke Vision as a tool

Advanced Training

YAML config file:

model: mlx-community/Qwen3-1.7B-4bit
data: ./training_data
iters: 1000
learning_rate: 1e-5
lora_parameters:
  rank: 8
  scale: 16.0
lr_schedule:
  name: cosine_decay
  warmup: 100
optimizer: adamw

Use config:

agent.tool.mlx_trainer(action="train", config="./lora_config.yaml")

Popular Models

Text:

  • mlx-community/Qwen3-1.7B-4bit (recommended)
  • mlx-community/Qwen3-4B-4bit
  • mlx-community/Llama-3.2-1B-4bit
  • mlx-community/gemma-2-2b-it-4bit

Vision:

  • mlx-community/Qwen2-VL-2B-Instruct-4bit (recommended)
  • mlx-community/Qwen2-Audio-7B-Instruct (audio)
  • mlx-community/llava-v1.6-mistral-7b-4bit

Community models at mlx-community


Troubleshooting

Out of memory:

config = {
    "grad_checkpoint": True,
    "batch_size": 1,
    "max_seq_length": 1024
}

Model degraded:

config = {
    "iters": 200,  # Lower for small datasets
    "learning_rate": 1e-5  # Conservative
}

Resources


Citation

@software{strands_mlx2025,
  author = {Cagatay Cali},
  title = {strands-mlx: MLX Model Provider for Strands Agents},
  year = {2025},
  url = {https://github.com/cagataycali/strands-mlx}
}

Apache 2 License | Built with MLX, MLX-LM, and Strands Agents

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

strands_mlx-0.2.5.tar.gz (42.4 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

strands_mlx-0.2.5-py3-none-any.whl (43.6 kB view details)

Uploaded Python 3

File details

Details for the file strands_mlx-0.2.5.tar.gz.

File metadata

  • Download URL: strands_mlx-0.2.5.tar.gz
  • Upload date:
  • Size: 42.4 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.6.5

File hashes

Hashes for strands_mlx-0.2.5.tar.gz
Algorithm Hash digest
SHA256 9b2900f0ce501e463a437777b28ef96bec7b51d4cb824353559b172d24ec0eff
MD5 45e8dbd91d2bcaaa9944a9e0c5a9b6c1
BLAKE2b-256 88bca2e51ed780b18e926eef57f12d17996459dfae501648604655d5d58dc20e

See more details on using hashes here.

File details

Details for the file strands_mlx-0.2.5-py3-none-any.whl.

File metadata

File hashes

Hashes for strands_mlx-0.2.5-py3-none-any.whl
Algorithm Hash digest
SHA256 92a8742720b03e5bbe971144752f439e2148976760d6635889394ef012bd49fd
MD5 76ad477510e71c12e506429e12436417
BLAKE2b-256 e883bba44773900c95ca058ff1799dca210de219aabeb1785f0da55805475da9

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page