Fine-Tuning Gemma 4 with Apple Silicon and MLX

Fine-tuning Gemma 4 on Google Colab is well-documented. But Colab has friction: free tier sessions time out, Pro sessions still have limits, and there's something uncomfortable about not knowing when your training run will get interrupted.

Running it on a Mac with Apple Silicon turns out to be a realistic alternative — at least for certain use cases. I tested this on an M3 Max with 96GB unified memory, and the experience was different from Colab in ways I didn't fully expect. Here's what I found. For an indie developer, having a training environment that adds zero recurring cloud cost is a quietly significant shift. In the second half, we'll connect the trained adapter all the way through Ollama into Antigravity.

Why MLX Over PyTorch + MPS

MLX is Apple's open-source machine learning framework designed specifically for Apple Silicon's unified memory architecture. Because the GPU and CPU share the same memory pool, a 96GB M3 Max can handle the full Gemma 4 27B model without quantization — barely, but it works.

PyTorch with MPS is an option too, but MLX proved more memory-efficient in practice, and the mlx-lm library provides fine-tuning utilities that integrate cleanly with the Hugging Face ecosystem.

Hardware note: if you're on an M1/M2/M3 chip with less than 64GB unified memory, use the 4-bit quantized version of Gemma 4 (approximately 14GB). Larger models on smaller machines are better handled by Colab.

Realistic Configurations by Unified Memory

What you can do locally is determined almost entirely by unified memory. Based on my own runs and configurations reported in the MLX community:

• 96GB+ (M3 Max / M4 Max high-end): 27B at 8-bit quantization + LoRA is comfortably stable. Start with --batch-size 4 and --lora-layers 16
• 64GB: 27B at 4-bit quantization + LoRA. Dropping to --batch-size 2 helps avoid swap
• 32-36GB: a 12B-class model at 4-bit is the realistic ceiling. 27B may run inference, but training gradients will exhaust memory
• 16-24GB: experiment with 4B-class models, or prepare datasets locally and hand training off to Colab

A rough planning formula that has served me well: quantized model size × 2 + 8GB equals your training memory budget. LoRA freezes the base model and trains only small adapter matrices, so it is far lighter than full fine-tuning — but gradients and optimizer state still add real overhead.

If you're wondering how much quality the quantization bit depth actually costs, Gemma 4 on Antigravity: Picking Q4 vs Q5 — What I Found After a Week on M2 Mac is a useful companion read. It compares inference rather than training, but the intuition about where quantization bites carries over.

Environment Setup

# Install Homebrew if not already installed
/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
 
# Python 3.11+ recommended
brew install python@3.11
 
# Create and activate virtual environment
python3.11 -m venv gemma4-mlx-env
source gemma4-mlx-env/bin/activate
 
# Install MLX and related libraries
pip install mlx-lm
pip install huggingface-hub datasets
 
# Log into Hugging Face (required to download Gemma 4)
huggingface-cli login

Before downloading Gemma 4, visit huggingface.co and accept the usage terms. Without this, the download returns a 401 error — easy to miss if you skip directly to the command.

Preparing the Dataset

Training data uses Alpaca format (instruction / input / output). For custom data, prepare a JSON file in the same structure:

# custom_dataset.json example
[
    {
        "instruction": "Find the bug in this code",
        "input": "def add(a, b):\n    return a - b",
        "output": "The operator is wrong. It should be `a + b`, not `a - b`."
    }
]

# Convert to MLX training format
import json
from pathlib import Path
 
def convert_to_mlx_format(input_file: str, output_dir: str):
    """Convert Alpaca-format JSON to MLX JSONL training format."""
    with open(input_file, "r", encoding="utf-8") as f:
        data = json.load(f)
    
    converted = []
    for item in data:
        prompt = f"<start_of_turn>user\n{item['instruction']}"
        if item.get("input"):
            prompt += f"\n\n{item['input']}"
        prompt += "<end_of_turn>\n<start_of_turn>model\n"
        response = item["output"] + "<end_of_turn>"
        converted.append({"text": prompt + response})
    
    # 90/10 train/validation split
    split = int(len(converted) * 0.9)
    
    out = Path(output_dir)
    out.mkdir(exist_ok=True)
    
    with open(out / "train.jsonl", "w") as f:
        for item in converted[:split]:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")
    
    with open(out / "valid.jsonl", "w") as f:
        for item in converted[split:]:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")
    
    print(f"Done: train={split}, valid={len(converted)-split}")
 
convert_to_mlx_format("custom_dataset.json", "data")

One mistake I made early on, so you can skip it: if you train on raw text without Gemma's chat template tags (<start_of_turn> / <end_of_turn>), loss goes down but conversations fall apart at inference time. When your base model is the instruction-tuned variant (-it), your training data must use the same template.

Running the LoRA Fine-Tune

# Download the model (first run only, ~14GB for 4-bit quantized version)
python -c "
from huggingface_hub import snapshot_download
snapshot_download('google/gemma-4-27b-it', local_dir='./gemma-4-27b-it')
"
 
# Start LoRA fine-tuning
mlx_lm.lora \
  --model ./gemma-4-27b-it \
  --train \
  --data ./data \
  --iters 1000 \
  --batch-size 4 \
  --lora-layers 16 \
  --learning-rate 1e-4 \
  --adapter-path ./adapters \
  --save-every 100

On M3 Max, 1,000 iterations took about 40 minutes. That's roughly 1.5-2x slower than an A100 on Colab Pro+, but the run completed without interruption — which matters more than you'd think when you're iterating on dataset quality.

With --save-every 100, adapter checkpoints are written every 100 iterations. For overnight runs this is essential: if the process dies, you can resume from the last checkpoint with --resume-adapter-file instead of starting over.

Memory and Time Budgets During Training (M3 Max / 96GB)

Even at the same 1,000 iterations, quantization and batch size change both the runtime and how much memory headroom you keep. Here are the approximate measured ranges from my own runs on an M3 Max (96GB). They're trends from a single machine, but they're enough for a first guess when picking a configuration.

The takeaway: doubling batch size does not halve the runtime. The gap between batch 4 and batch 2 at 8-bit stays around 1.4x, because on Apple Silicon memory bandwidth tends to be the bottleneck — larger batches don't pack the compute linearly. Flip that around and dropping to batch 2 when memory is tight costs less time than you'd fear. Prioritize avoiding swap first; you'll usually finish sooner that way.

Choosing a Realistic Learning Rate and Iteration Count

Treat the sample command's --learning-rate 1e-4 as a value tuned for datasets in the thousands of examples. When I ran a few-hundred-example dataset at 1e-4, loss dropped fast — too fast. The model slid into memorization, parroting phrasings straight out of the training data. Dropping to 5e-5 kept the responses natural while still pulling them toward the target style. My rule now: start at 5e-5, and only raise it if loss barely moves.

Organizing the behavior by learning rate, from what I observed on the same 300-example dataset:

For iteration count, my starting point is "dataset size × 2-3 epochs ÷ batch size." With 300 examples and --batch-size 4, one epoch is roughly 75 iterations, so 150-225 iterations is a sensible first run. There's no need to default to 1,000.

Validation loss is the signal that matters. When the valid loss printed at each checkpoint stops improving, further iterations only deepen memorization. Watching train loss alone is the trap — it keeps "improving" convincingly, and I wasted several early runs exactly this way.

Choosing LoRA Rank and Scale

--lora-layers controls how many layers get an adapter attached. What determines the adapter's actual capacity is the rank — and if you only ever run from the command line, that rank quietly stays at its default. I missed this for longer than I'd like to admit.

MLX lets you move the finer hyperparameters into a YAML config. Once you start running comparisons, this is worth doing purely for reproducibility.

# lora_config.yaml
model: "./gemma-4-27b-it"
train: true
data: "./data"
adapter_path: "./adapters"
iters: 300
batch_size: 2
learning_rate: 5e-5
save_every: 100
lora_parameters:
  keys: ["self_attn.q_proj", "self_attn.v_proj"]
  rank: 16
  scale: 20.0
  dropout: 0.05

# Run from the config file (command-line flags override config values)
caffeinate -i mlx_lm.lora --config lora_config.yaml 2>&1 | tee train_log.txt

One detail trips up anyone arriving from another framework. Hugging Face PEFT asks for lora_alpha; MLX's lora_parameters asks for scale. It isn't just a rename — PEFT divides alpha by rank internally, while MLX applies scale directly as a multiplier. Copy a PEFT config over verbatim and you can end up training several times stronger than intended. The safe starting translation is scale ≈ lora_alpha / rank.

What rank to pick depends on what you're actually teaching the model. Holding the dataset fixed at 300 examples and varying only rank, the pattern looked like this:

With 300 examples at rank 32, validation loss turned upward around iteration 150. The extra capacity simply memorized the data faster than it generalized. While your dataset is still in the hundreds, stay at rank 8–16 and grow the data before you grow the rank. That order wastes fewer nights.

The keys list matters too. The default q_proj / v_proj pair is the lightest configuration and is enough when you're shaping output behavior. Adding entries like mlp.gate_proj helps when you're pushing knowledge in, but memory use and runtime both climb accordingly. Get one full cycle working on the defaults first — otherwise you can't tell which change produced the result.

Setting Up Your Mac for Overnight Runs

The Mac's biggest advantage over Colab is running while you sleep. But run unprepared and you'll wake up to a machine that went to sleep mid-training. This is the setup I settled on:

# Block idle sleep for exactly as long as training runs
caffeinate -i mlx_lm.lora \
  --model ./gemma-4-27b-it \
  --train \
  --data ./data \
  --iters 1000 \
  --batch-size 4 \
  --lora-layers 16 \
  --learning-rate 5e-5 \
  --adapter-path ./adapters \
  --save-every 100 2>&1 | tee train_log.txt

• caffeinate -i blocks idle sleep only while the command is alive. No system settings to change, nothing to remember to switch back. The display can sleep — that's fine. You only need to stop system sleep
• Plug into power. On battery, macOS leans toward power saving and training speed becomes unstable
• tee writes the log to a file, so in the morning you can review the full loss curve. Relying on terminal scrollback alone means losing the middle of the run
• Temporarily pause macOS automatic updates that involve restarts, at least on nights you train

Heat and Effective Throughput on Long Runs

The surprise from overnight runs was that speed drifts downward over time. Comparing iters/sec at the ten-minute mark against the three-hour mark on the same M3 Max configuration, throughput was down by just under 10 percent. That sounds trivial, but it means extrapolating total runtime from your first 100 iterations lands 5–10% optimistic. If a run you expected to finish overnight is still going at breakfast, this is usually why.

You can watch the thermal picture with a second terminal open alongside the run:

# Sample SMC sensors every 5 seconds (Ctrl-C to stop)
sudo powermetrics --samplers smc -i 5000 | grep -i -E "temp|fan"

The countermeasures are unglamorous. On a laptop, leaving the lid open with airflow underneath held up better than running clamshell into an external display. Keep the power adapter connected and close anything else heavy. And keep --save-every in the command — if thermals do destabilize the run, everything up to the last checkpoint survives.

Resuming is a matter of pointing at that checkpoint:

# Resume training from the most recent adapter
caffeinate -i mlx_lm.lora \
  --model ./gemma-4-27b-it \
  --train \
  --data ./data \
  --iters 400 \
  --batch-size 2 \
  --lora-layers 16 \
  --learning-rate 5e-5 \
  --adapter-path ./adapters \
  --resume-adapter-file ./adapters/adapters.safetensors

One detail worth internalizing: --iters after a resume means additional iterations, not a cumulative target. I misread that once and ran twice the training I intended.

Testing the Fine-Tuned Model

from mlx_lm import load, generate
 
model, tokenizer = load(
    "./gemma-4-27b-it",
    adapter_path="./adapters"
)
 
prompt = """<start_of_turn>user
Find any issues with this Python code:
 
def calculate_average(numbers):
    total = sum(numbers)
    return total / len(numbers)
<end_of_turn>
<start_of_turn>model
"""
 
response = generate(model, tokenizer, prompt=prompt, max_tokens=512, temp=0.7)
print(response)

A practical testing tip: prompt the model with paraphrases of your training questions, not the originals. Answering memorized inputs proves nothing — generalization shows when the response style holds up under rewording. Small datasets of a few hundred examples overfit quickly; if validation loss starts climbing mid-run, reduce --iters or add data.

Measuring Whether the Fine-Tune Actually Helped

Running a handful of prompts and concluding "that feels better" is the easiest way to fool yourself on a small fine-tune. Without holding conditions constant, you cannot tell whether a difference came from the adapter or from sampling noise.

What I do instead is keep around ten prompts that never appeared in training, and write both models' answers side by side:

# eval_ab.py — same prompts, base model vs. fine-tuned
import json
from mlx_lm import load, generate
 
HELD_OUT = [
    "Explain why this function returns the wrong value:\n\ndef ratio(a, b):\n    return a / b",
    "Rewrite this loop as a list comprehension:\n\nout = []\nfor x in xs:\n    out.append(x * 2)",
]
 
def build_prompt(question: str) -> str:
    return f"<start_of_turn>user\n{question}<end_of_turn>\n<start_of_turn>model\n"
 
def run(model_path: str, adapter_path):
    model, tokenizer = load(model_path, adapter_path=adapter_path)
    return [
        # temp=0.0 keeps the comparison from drowning in sampling noise
        generate(model, tokenizer, prompt=build_prompt(q), max_tokens=384, temp=0.0)
        for q in HELD_OUT
    ]
 
base = run("./gemma-4-27b-it", None)
tuned = run("./gemma-4-27b-it", "./adapters")
 
rows = [
    {"prompt": q, "base": b, "tuned": t, "base_len": len(b), "tuned_len": len(t)}
    for q, b, t in zip(HELD_OUT, base, tuned)
]
 
with open("ab_result.json", "w", encoding="utf-8") as f:
    json.dump(rows, f, ensure_ascii=False, indent=2)
 
for r in rows:
    print(f"{r['prompt'][:20]}... base={r['base_len']} chars / tuned={r['tuned_len']} chars")

Pinning temperature to 0.0 is the part that matters. Saving the output to JSON also means you can reread it with fresh eyes the next morning, and diff one training configuration against another.

Three things are worth reading for:

The third row comes from experience. When responses stopped terminating, I assumed overfitting; the actual cause was a branch in my conversion script that dropped the closing tag on a subset of records. Catching that after a full training run is expensive, so eyeballing the last few lines of train.jsonl is time well spent.

Keeping Several Adapters Instead of Fusing

Fusing writes out a complete model, so each fuse of a 27B base costs tens of gigabytes. Do it once per use case and your disk disappears.

What settled into place for me was fusing only right before distribution or Ollama registration, and otherwise swapping adapters in place. A rank-16 adapter is only 10–20MB, so keeping one per task costs essentially nothing.

adapters/
├── code-review/      # Review-comment style
├── commit-message/   # Commit message conventions
└── doc-ja/           # Japanese documentation voice

Switching means changing adapter_path at load time. The base model stays a single copy on disk.

# switch_adapter.py — swap adapters against one shared base model
import sys
from mlx_lm import load, generate
 
ADAPTERS = {
    "code-review": "./adapters/code-review",
    "commit-message": "./adapters/commit-message",
    "doc-ja": "./adapters/doc-ja",
    "base": None,  # no adapter, for comparison
}
 
name = sys.argv[1] if len(sys.argv) > 1 else "base"
question = sys.argv[2] if len(sys.argv) > 2 else "Write a commit message for this diff"
 
model, tokenizer = load("./gemma-4-27b-it", adapter_path=ADAPTERS[name])
prompt = f"<start_of_turn>user\n{question}<end_of_turn>\n<start_of_turn>model\n"
print(generate(model, tokenizer, prompt=prompt, max_tokens=384, temp=0.0))

python switch_adapter.py code-review "What's wrong with this code?"
python switch_adapter.py base "What's wrong with this code?"

Because the same question can be fired at any adapter by name, the evaluation step from the previous section becomes something you do casually rather than as a separate ritual. Save each run to its own directory and you can line up generations — rank16-300iter next to rank8-500iter — and compare them directly.

Two caveats. An adapter is only valid against the base it was trained on, so swapping an 8-bit base for a 4-bit one fails at load time; putting the quantization in the directory name saves you from that. And adapters don't stack — if you want both a voice and a body of knowledge, you mix them in the training data rather than layering two adapters.

Fusing the Adapter for Ollama / Antigravity

Once you are happy with the adapter, fuse it into a standalone model:

mlx_lm.fuse \
  --model ./gemma-4-27b-it \
  --adapter-path ./adapters \
  --save-path ./gemma-4-27b-custom

The fused model works directly with mlx_lm.generate, and if you convert it to GGUF with llama.cpp's conversion script, it loads into Ollama. From there, Antigravity's local LLM connection can call it like any other model — meaning a completion model trained on your own codebase's conventions can live inside your IDE, all built on a single Mac. The Antigravity-side connection steps are covered in Setting Up Local LLMs in Antigravity for Practical Use.

The GGUF step uses llama.cpp's conversion script against the fused model. It is short, but two details will save you a wasted evening.

# Grab llama.cpp (only the conversion tooling is needed here)
git clone --depth 1 https://github.com/ggml-org/llama.cpp
pip install -r llama.cpp/requirements.txt
 
# Convert the fused model to f16 GGUF
python llama.cpp/convert_hf_to_gguf.py ./gemma-4-27b-custom \
  --outfile ./gemma-4-27b-custom-f16.gguf \
  --outtype f16
 
# Quantize to a practical size (Q4_K_M balances quality and footprint well)
llama.cpp/build/bin/llama-quantize \
  ./gemma-4-27b-custom-f16.gguf \
  ./gemma-4-27b-custom-q4_k_m.gguf Q4_K_M

First, if you fine-tuned against a quantized base, mlx_lm.fuse needs --de-quantize or the conversion script will choke on weights it cannot interpret. Second, watch disk: the f16 GGUF for a 27B model runs past 50GB, so budget roughly 70GB free including the quantized output. The f16 intermediate can be deleted once quantization finishes.

# After GGUF conversion, register with Ollama
cat > Modelfile <<'EOF2'
FROM ./gemma-4-27b-custom.gguf
PARAMETER temperature 0.7
EOF2
ollama create gemma4-custom -f Modelfile
ollama run gemma4-custom

One more thing worth bookmarking: if long responses start getting cut off once you're running through Ollama, Fixing Mid-Stream Cutoffs and Long-Run Freezes When Antigravity Talks to Ollama covers the fixes. The symptom isn't specific to custom models — it affects Ollama integrations in general.

Common Errors and Fixes

Errors I actually hit, or reproduced while testing:

• Out of memory right after training starts: halve --batch-size first. If it still dies, drop --lora-layers from 16 to 8 or switch to the 4-bit base model. Sometimes just closing the browser is enough. If the same class of error appears later on the inference side — when calling the model from Antigravity — see How to Fix Out of Memory Errors When Using Gemma 4 in Antigravity
• 401 Unauthorized on download: you have not accepted the Gemma license on Hugging Face, or your CLI token lacks read permission
• Training progresses but responses are broken: the chat template mismatch described above. Pull one line out of train.jsonl and visually confirm the template tags are present
• Sudden slowdown mid-run: memory pressure has pushed you into swap. If Activity Monitor's memory pressure turns yellow, a smaller batch size finishes faster overall

When to Use Mac vs. Colab

The choice comes down to dataset size and model scale:

Use Mac (MLX) when:

• Dataset is under ~1,000 examples (fast iteration, low stakes)
• You need overnight or multi-hour runs without session limits
• You want to avoid cloud costs entirely
• You have 64GB+ unified memory

Use Colab when:

• Dataset is 10,000+ examples and training time matters
• You're working with models larger than Gemma 4 27B
• Your Mac has less than 32GB unified memory

The practical pattern I've settled on: prototype and iterate on Mac, then run the final production training pass on Colab Pro+ where speed matters. MLX made the Mac genuinely useful for fine-tuning — not just as an inference machine, but as a real part of the training workflow.

Start with a small dataset of around 100 examples and push one full cycle through the pipeline — convert, train, fuse, register with Ollama. Once the pipeline works end to end, you can spend your energy where it matters: improving the data. I hope this helps if you are building on the same setup.