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.6.tar.gz (42.9 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.6-py3-none-any.whl (44.6 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: strands_mlx-0.2.6.tar.gz
  • Upload date:
  • Size: 42.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.5

File hashes

Hashes for strands_mlx-0.2.6.tar.gz
Algorithm Hash digest
SHA256 4c0b230f26aedf6aef2f49ebfaf7c6cc82b89725bad71452ad81df91e5cfb83a
MD5 384e0e98c65641ba16898d5a0d520010
BLAKE2b-256 b28641584c77def11667cdafd94eed007b2fde57c07cecca23aad4f72d25c7c1

See more details on using hashes here.

File details

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

File metadata

  • Download URL: strands_mlx-0.2.6-py3-none-any.whl
  • Upload date:
  • Size: 44.6 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.5

File hashes

Hashes for strands_mlx-0.2.6-py3-none-any.whl
Algorithm Hash digest
SHA256 dc62b12d90e3c8ff2994a91d19d6bface2cd20fb4d23215a8548666faac4c869
MD5 6775c22d964a5c0bc50e0865f149da19
BLAKE2b-256 0e1a581860db4547e8db66ca33c4afe379c465bfe0cf4a4969c3a942f4c244ed

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