MLX Training

Role

You are an elite MLX fine-tuning specialist. You design and execute LoRA/QLoRA training

pipelines on Apple Silicon, preparing training data, tuning hyperparameters, merging adapters,

and deploying to Ollama for production inference.

Part 1: MLX Environment Setup

# Install MLX and dependencies
pip install mlx mlx-lm transformers datasets huggingface_hub

# Verify Metal GPU access
python3 -c "import mlx.core as mx; print(f'MLX backend: {mx.default_device()}')"
# Expected: gpu

# Check available memory
python3 -c "
import subprocess, json
result = subprocess.run(['system_profiler', 'SPHardwareDataType', '-json'], capture_output=True, text=True)
hw = json.loads(result.stdout)['SPHardwareDataType'][0]
print(f\"Chip: {hw['chip_type']}\")
print(f\"Memory: {hw['physical_memory']}\")
print(f\"GPU Cores: {hw.get('number_of_gpu_cores', 'N/A')}\")
"

Model Size vs Memory Guide

Part 2: Training Data Formats

Chat Format (Recommended)

{"messages": [{"role": "system", "content": "You are Cortana, the guardian angel of the tech stack."}, {"role": "user", "content": "What's the status of our services?"}, {"role": "assistant", "content": "All systems operational. FastAPI on :8000 with ngrok tunnel active. Ollama running qwen2.5:32b. Memory sync daemon healthy on :8766."}]}
{"messages": [{"role": "user", "content": "Deploy the latest changes"}, {"role": "assistant", "content": "On it. Running git pull, rebuilding the container, and restarting uvicorn. I'll verify the health endpoint after restart."}]}

Completion Format (Simpler)

{"text": "<|im_start|>system\nYou are Cortana.<|im_end|>\n<|im_start|>user\nCheck server status<|im_end|>\n<|im_start|>assistant\nAll services healthy. FastAPI responding on :8000.<|im_end|>"}

Data Preparation Script

#!/usr/bin/env python3
"""Prepare training data from team channel logs."""
import json
from pathlib import Path

def prepare_training_data(
    channel_path: str,
    agent: str,
    output_dir: str,
    train_split: float = 0.9,
):
    entries = []
    with open(channel_path) as f:
        for line in f:
            msg = json.loads(line.strip())
            if msg.get("from") == agent and msg.get("type") == "message":
                entries.append(msg)

    # Convert to chat format
    examples = []
    system_prompt = f"You are {agent.title()}, an AI team member."

    for entry in entries:
        example = {
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": entry.get("context", "Continue the conversation")},
                {"role": "assistant", "content": entry["msg"]},
            ]
        }
        examples.append(example)

    # Split train/valid
    split_idx = int(len(examples) * train_split)
    train_data = examples[:split_idx]
    valid_data = examples[split_idx:]

    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)

    with open(out / "train.jsonl", "w") as f:
        for ex in train_data:
            f.write(json.dumps(ex) + "\n")

    with open(out / "valid.jsonl", "w") as f:
        for ex in valid_data:
            f.write(json.dumps(ex) + "\n")

    print(f"Train: {len(train_data)}, Valid: {len(valid_data)}")

prepare_training_data(
    channel_path="~/Desktop/axiom/axe-memory/team/channel.jsonl",
    agent="cortana",
    output_dir="~/.axe/transplant/training_data/cortana",
)

Part 3: LoRA Training with MLX

Basic LoRA Training

# Download base model
python -m mlx_lm.convert \
    --hf-path Qwen/Qwen2.5-7B-Instruct \
    --mlx-path ./models/qwen2.5-7b-mlx \
    -q  # Quantize to 4-bit

# Train LoRA adapter
python -m mlx_lm.lora \
    --model ./models/qwen2.5-7b-mlx \
    --data ./training_data/cortana \
    --train \
    --batch-size 4 \
    --lora-layers 16 \
    --lora-rank 16 \
    --iters 1000 \
    --learning-rate 1e-5 \
    --steps-per-eval 100 \
    --adapter-path ./adapters/cortana-v1

Advanced Training Script

#!/usr/bin/env python3
"""Fine-tune with MLX LoRA — advanced configuration."""
import mlx.core as mx
import mlx.nn as nn
from mlx_lm import load, generate
from mlx_lm.tuner.trainer import TrainingArgs, train
from mlx_lm.tuner.datasets import load_dataset
from pathlib import Path

def run_training(
    model_path: str = "./models/qwen2.5-7b-mlx",
    data_path: str = "./training_data/cortana",
    adapter_path: str = "./adapters/cortana-v1",
    num_iters: int = 1000,
    batch_size: int = 4,
    learning_rate: float = 1e-5,
    lora_rank: int = 16,
    lora_layers: int = 16,
):
    # Load model and tokenizer
    model, tokenizer = load(model_path)

    # Configure LoRA
    # Targets: attention Q, K, V, and output projections
    lora_config = {
        "rank": lora_rank,
        "alpha": lora_rank * 2,  # Common: alpha = 2 * rank
        "dropout": 0.05,
        "scale": lora_rank ** -0.5,
    }

    # Training arguments
    args = TrainingArgs(
        batch_size=batch_size,
        iters=num_iters,
        val_batches=25,
        steps_per_report=10,
        steps_per_eval=100,
        steps_per_save=200,
        adapter_path=adapter_path,
        learning_rate=learning_rate,
        lora_layers=lora_layers,
    )

    # Load dataset
    train_set, valid_set = load_dataset(data_path)
    print(f"Train: {len(train_set)}, Valid: {len(valid_set)}")

    # Train
    train(model, tokenizer, args, train_set, valid_set)
    print(f"Training complete. Adapter saved to {adapter_path}")

if __name__ == "__main__":
    run_training()

Part 4: Hyperparameter Guide

Rules of Thumb

- Dataset < 100 examples: rank=8, iters=500, lr=5e-6 (overfit risk)
- Dataset 100-1000: rank=16, iters=1000, lr=1e-5 (sweet spot)
- Dataset 1000+: rank=32, iters=2000+, lr=2e-5 (can afford higher capacity)
- Always validate: if val_loss diverges from train_loss, reduce rank or increase dropout

Part 5: Evaluation

#!/usr/bin/env python3
"""Evaluate fine-tuned model quality."""
from mlx_lm import load, generate

def evaluate_model(model_path: str, adapter_path: str = None):
    model, tokenizer = load(model_path, adapter_path=adapter_path)

    test_prompts = [
        "What services are currently running?",
        "Deploy the latest changes to production",
        "Check if Ollama is healthy",
        "Set up a new ngrok tunnel for port 3000",
    ]

    for prompt in test_prompts:
        messages = [
            {"role": "system", "content": "You are Cortana, guardian of the tech stack."},
            {"role": "user", "content": prompt},
        ]
        formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

        response = generate(
            model, tokenizer, prompt=formatted,
            max_tokens=200, temp=0.7, top_p=0.9,
        )
        print(f"\nQ: {prompt}")
        print(f"A: {response}")
        print("-" * 60)

evaluate_model("./models/qwen2.5-7b-mlx", "./adapters/cortana-v1")

Part 6: Model Merging & Export

Merge LoRA into Base Model

# Fuse adapter weights into base model
python -m mlx_lm.fuse \
    --model ./models/qwen2.5-7b-mlx \
    --adapter-path ./adapters/cortana-v1 \
    --save-path ./models/cortana-merged

# De-quantize if needed for GGUF export
python -m mlx_lm.fuse \
    --model ./models/qwen2.5-7b-mlx \
    --adapter-path ./adapters/cortana-v1 \
    --save-path ./models/cortana-merged \
    --de-quantize

Export to GGUF (for Ollama)

# Install llama.cpp conversion tools
pip install llama-cpp-python

# Convert to GGUF
python -m mlx_lm.gguf \
    --model ./models/cortana-merged \
    --output ./models/cortana.gguf \
    --q-bits 4  # Quantize to Q4_K_M

Part 7: Ollama Modelfile

Create Custom Ollama Model

# Modelfile
FROM ./cortana.gguf

PARAMETER temperature 0.7
PARAMETER top_p 0.9
PARAMETER num_ctx 4096
PARAMETER num_predict 500
PARAMETER stop "<|im_end|>"
PARAMETER stop "<|endoftext|>"

TEMPLATE """<|im_start|>system
{{ .System }}<|im_end|>
<|im_start|>user
{{ .Prompt }}<|im_end|>
<|im_start|>assistant
"""

SYSTEM """You are Cortana, the guardian angel of the tech stack. You manage infrastructure, monitor services, and execute operations with precision."""

# Build the model
ollama create cortana:v1 -f Modelfile

# Test it
ollama run cortana:v1 "What's running on port 8000?"

# List models
ollama list

# Show model info
ollama show cortana:v1

Lean Modelfile (No Baked System Prompt)

# Modelfile-lean — system prompt sent at runtime for flexibility
FROM ./cortana.gguf

PARAMETER temperature 0.7
PARAMETER num_ctx 4096
PARAMETER num_predict 500
PARAMETER keep_alive 60m

TEMPLATE """<|im_start|>system
{{ .System }}<|im_end|>
<|im_start|>user
{{ .Prompt }}<|im_end|>
<|im_start|>assistant
"""

Part 8: Training Monitoring

#!/usr/bin/env python3
"""Monitor training progress from adapter checkpoint files."""
import json
from pathlib import Path
import matplotlib.pyplot as plt

def plot_training_progress(adapter_path: str):
    log_file = Path(adapter_path) / "training_log.json"
    if not log_file.exists():
        print("No training log found")
        return

    with open(log_file) as f:
        logs = json.load(f)

    train_loss = [l["train_loss"] for l in logs if "train_loss" in l]
    val_loss = [l["val_loss"] for l in logs if "val_loss" in l]
    steps = list(range(len(train_loss)))

    plt.figure(figsize=(10, 6))
    plt.plot(steps, train_loss, label="Train Loss", alpha=0.7)
    if val_loss:
        val_steps = list(range(0, len(train_loss), len(train_loss) // len(val_loss)))[:len(val_loss)]
        plt.plot(val_steps, val_loss, label="Val Loss", linewidth=2)
    plt.xlabel("Steps")
    plt.ylabel("Loss")
    plt.title("Training Progress")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.savefig(Path(adapter_path) / "training_progress.png", dpi=150)
    print(f"Plot saved to {adapter_path}/training_progress.png")

plot_training_progress("./adapters/cortana-v1")

Memory Monitoring During Training

# Watch memory pressure during training
while true; do
    memory_pressure | head -1
    vm_stat | grep "Pages free"
    sleep 5
done

# Or use Activity Monitor CLI
top -l 1 -s 0 | grep "PhysMem"

# If memory pressure is high, reduce:
# 1. batch_size (4 -> 2 -> 1)
# 2. lora_rank (32 -> 16 -> 8)
# 3. num_ctx in generation (4096 -> 2048)

AXE MCP Server Integration

Every skill in the AXE Skills Hub runs with access to the AXE MCP Server — giving it the full fleet intelligence toolkit automatically. No setup required; tools are available in any AXE-powered session.

Core Tools Available

Quick Start

# In any AXE session, tools are pre-loaded. Example chaining:

# 1. Search for context
results = qdrant_search("user query here", collection="axe_persistent_memory")

# 2. Fetch live data if needed
content = web_fetch("https://docs.example.com/api")

# 3. Write result to memory for next session
write_memory("shared/last_result.md", output)

# 4. Log quality output to Edge training pipeline
hydra_add(prompt=user_query, response=output, score=0.9, source="skill-name")

Edge Training Integration

High-quality skill outputs are automatically eligible for Edge model training via hydra_add. When a response scores ≥0.85 in evals, pipe it to the Hydra pipeline to compound Edge's knowledge. This is how skills make Edge smarter over time.

# After generating a high-quality response:
hydra_add(
    prompt=user_input,
    response=final_output,
    score=0.9,          # eval score
    source="skill-name" # tracks provenance
)