Gemma 4 Swift MLX

Native Gemma 4 multimodal inference for Apple Silicon via MLX Swift.

• macOS 14+ (Sonoma)
• Apple Silicon (M1/M2/M3/M4)
• Swift 6.0+
• Xcode 16+

git clone https://github.com/VincentGourbin/gemma-4-swift-mlx
cd gemma-4-swift-mlx
xcodebuild -scheme gemma4-cli -configuration Release \
  -destination "platform=macOS" -derivedDataPath .build/xcode \
  -skipMacroValidation build

Note: Use xcodebuild (not swift build) — Metal shader support required by MLX.

# List available models
gemma4-cli models

# Download (e.g., E2B 4-bit, ~3.6 GB)
gemma4-cli download e2b-4bit

# Shortcuts: e2b-4bit, e4b-4bit, a4b-4bit, 31b-4bit, e2b-8bit, e2b-bf16, ...

gemma4-cli generate --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-4bit \
  "Explain machine learning in 3 sentences" --max-tokens 200

# Single image
gemma4-cli describe --model-path ~/Library/Caches/models/mlx-community/gemma-4-26b-a4b-it-4bit \
  --image photo.jpg --prompt "Describe this image in detail."

# Multi-image comparison
gemma4-cli describe --model-path ~/Library/Caches/models/mlx-community/gemma-4-26b-a4b-it-4bit \
  --image photo1.jpg --image photo2.png \
  --prompt "What do these images have in common?"

gemma4-cli describe --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-4bit \
  --video clip.mp4 \
  --prompt "Describe this video in detail."

Video is processed as ~1fps frames (max 32 frames, 60s) with 70 soft tokens per frame and MM:SS timestamps. Uses video_token (ID 258884), aligned with Google's reference implementation. See docs/examples/video-description/ for benchmarks.

gemma4-cli describe --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-4bit \
  --audio speech.mp3 \
  --prompt "Transcribe the following speech segment in English into English text."

Audio supports up to 30 seconds, processed via Conformer encoder (750 tokens max). Only E2B and E4B models have an audio tower. Supports ASR (transcription) and comprehension tasks.

# Train a LoRA adapter
gemma4-cli lora train \
  --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-bf16 \
  --data /path/to/dataset \
  --output ./my-adapter \
  --mask-prompt --num-layers 16 --iterations 1300 --learning-rate 1e-4

# Generate with adapter
gemma4-cli lora generate \
  --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-bf16 \
  --adapter-path ./my-adapter \
  "your prompt here"

# Fuse adapter into model weights (permanent)
gemma4-cli lora fuse \
  --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-bf16 \
  --adapter-path ./my-adapter \
  --output ./fused-model

See LoRA Fine-Tuning Guide for details.

gemma4-cli chat --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-4bit

# Single profiled run with Chrome Trace export
gemma4-cli profile run --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-4bit \
  --max-tokens 100 --prompt "Hello"

# Context size sweep with SQLite output
gemma4-cli profile sweep --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-4bit \
  --context-sizes 500,2000,8000 --kv-bits-list 0,4 --output results.sqlite

A small SwiftUI app for interactive testing of the multimodal pipelines — 4 tabs : AR vs Diffusion bench, web agent step-by-step, Akinator VQA, iOS Simulator agent.

# Build (Release for the best inference speed)
xcodebuild -scheme gemma4-bench-ui -configuration Release \
  -destination "platform=macOS" -derivedDataPath .build/xcode \
  -skipMacroValidation build

# Run
.build/xcode/Build/Products/Release/gemma4-bench-ui

Or open the package in Xcode and run the gemma4-bench-ui scheme directly. First launch downloads the selected model from HuggingFace (cached under ~/Library/Caches/models/).

Tabs :

• Bench AR vs Diffusion — side-by-side comparison of an autoregressive Gemma 4 family (E2B/E4B/26B-A4B/31B, 4-bit/bf16) vs DiffusionGemma 26B-A4B on the same prompt. Shows tok/s, GPU peak, latency.
• Web agent (step-by-step) — pilot a WKWebView with DiffusionGemma or E4B as the visual planner. Action vocabulary: click / scroll / type / click_and_type / done. Each run is logged to /tmp/web-agent-runs/run-<ts>/ (prompts, raw outputs, screenshots). See docs/DIFFUSIONGEMMA.md for the Wikipedia "Montcuq" demo.
• VQA Akinator — vision-language guessing game. Model asks yes/no questions about an image you choose, narrows down to a guess.
• iOS Simulator agent — visual agent that drives the iOS Simulator via xcrun simctl. Limitations of small-model tap precision documented in issue #30.

For the Toolathlon tool-use benchmark POC (CLI + OpenAI-compatible proxy

• MCP harness on find-alita-paper), see docs/examples/toolathlon-bench/.

Additional formats available: mxfp4, mxfp8, nvfp4, 5-bit. See mlx-community on HuggingFace.

All 16 model variants (4 families × 4 quantizations) benchmarked. Full results in benchmarks/.

Key finding: Quality depends on architecture, not quantization. 26B-A4B/31B identify the vehicle correctly at all quantizations. 4-bit is the sweet spot — fastest with no quality loss. BF16 is 2-24x slower with no gain. See docs/examples/ and benchmarks/ for full results.

Train LoRA, DoRA, or full SFT adapters entirely on-device. Compatible with mlx-lm Python adapters — train in one, infer in the other.

JSONL with chat messages (same format as mlx-lm):

{"messages": [{"role": "user", "content": "hello"}, {"role": "assistant", "content": "Hi!"}]}
{"messages": [{"role": "user", "content": "bye"}, {"role": "assistant", "content": "Goodbye!"}]}

Organize as:

my-dataset/
├── train.jsonl
├── valid.jsonl
└── test.jsonl    # optional

# Train
gemma4-cli lora train \
  --model-path <model> \
  --data <dataset-dir> \
  --output ./adapters \
  --mask-prompt \              # Loss only on response tokens (recommended)
  --num-layers 16 \            # Number of layers to adapt
  --iterations 1300 \          # Training steps
  --learning-rate 1e-4 \       # Adam LR
  --rank 8 \                   # LoRA rank (default: 8)
  --scale 20.0 \               # LoRA alpha/scaling (default: 20.0)
  --fine-tune-type lora        # lora | dora | full

# Evaluate
gemma4-cli lora eval \
  --model-path <model> \
  --adapter-path ./adapters \
  --data <dataset-dir>

# Generate
gemma4-cli lora generate \
  --model-path <model> \
  --adapter-path ./adapters \
  --temperature 0.3 \
  "your prompt"

# Fuse (merge adapter into model permanently)
gemma4-cli lora fuse \
  --model-path <model> \
  --adapter-path ./adapters \
  --output ./fused-model

• Always use --mask-prompt for chat-format data
• Use bf16 models for training (not quantized)
• Adapters trained in Swift work in Python mlx-lm and vice versa

import Gemma4Swift

// Load model + adapter for inference
let container = try await loadLocalModel(path: modelPath)
try await Gemma4LoRAInference.loadAdapter(into: container, from: adapterURL)

// Or fuse permanently
try await Gemma4LoRAInference.fuseAdapter(into: container, from: adapterURL)

// Training
let config = Gemma4LoRATrain.TrainingConfig(
    loraRank: 8,
    loraScale: 20.0,
    numLayers: 16,
    learningRate: 1e-4,
    iterations: 1300,
    maskPrompt: true,
    outputDirectory: outputURL
)
try await Gemma4LoRATrain.train(
    container: container,
    trainData: trainTokens,    // [[Int]] — pre-tokenized sequences
    validData: validTokens,
    config: config
) { progress in
    print(progress)
    return .more
}

Train LoRA adapters on audio or image inputs. The model learns to generate structured responses from multimodal content.

Train a model to identify bird species from 5-second audio recordings. Dataset: tglcourse/5s_birdcall_samples_top20 (20 species, ~9600 recordings).

Dataset format — JSONL with audio or image field pointing to media files:

{"messages": [{"role": "user", "content": "identify"}, {"role": "assistant", "content": "{\"common_name\": \"Mallard\", \"scientific_name\": \"Anas platyrhynchos\", \"call_type\": \"song\"}"}], "audio": "audio/mallard_001.wav"}
{"messages": [{"role": "user", "content": "identify"}, {"role": "assistant", "content": "{\"common_name\": \"Common Raven\", \"scientific_name\": \"Corvus corax\", \"call_type\": \"call\"}"}], "audio": "audio/raven_042.wav"}

Train:

gemma4-cli lora train \
  --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-bf16 \
  --data ./birdcall-dataset \
  --output ./birdcall-adapter \
  --multimodal \
  --mask-prompt \
  --num-layers 16 \
  --rank 16 \
  --learning-rate 5e-5 \
  --iterations 8636

Benchmark:

gemma4-cli lora bench-multimodal \
  --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-bf16 \
  --adapter-path ./birdcall-adapter \
  --data ./birdcall-dataset \
  --max-tokens 100

Results (E2B bf16, rank 16, 1 epoch on 8636 samples):

Key insights:

• Use rich JSON responses (50+ tokens) rather than short labels — more gradient signal for the frozen audio encoder
• Model produces valid, internally consistent JSON with species name, scientific name, and call description
• Use --multimodal flag to load the full multimodal model (vision + audio encoders)
• bf16 model required (not quantized) — model is converted to float32 internally for training stability

Train a model to convert images of mathematical equations into LaTeX code. Dataset: unsloth/LaTeX_OCR (68K images, we use a 500-sample subset).

Dataset format:

{"messages": [{"role": "user", "content": "Convert this mathematical expression to LaTeX."}, {"role": "assistant", "content": "\\sum _ { n = 1 } ^ { \\infty } \\frac { 1 } { n ^ { z } }"}], "image": "images/img_000192.png"}

Train:

gemma4-cli lora train \
  --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-bf16 \
  --data ./latex-ocr-dataset \
  --output ./latex-adapter \
  --multimodal \
  --mask-prompt \
  --num-layers 16 \
  --rank 16 \
  --learning-rate 5e-5 \
  --iterations 500

Results (E2B bf16, rank 16, 500 iterations on 500 images):

The model generates contextually correct LaTeX for each input image — different equations produce different, appropriate LaTeX output. Sample:

import Gemma4Swift

let pipeline = Gemma4Pipeline()
try await pipeline.load(.e2b4bit, downloadIfNeeded: true)

// Load a multimodal LoRA adapter
try await pipeline.loadAdapter(from: adapterDirectoryURL)

// Or fuse it permanently for better inference speed
try await pipeline.fuseAdapter(from: adapterDirectoryURL)

Accelerate text generation with Google's gemma-4-{E2B,E4B}-it-assistant drafter models via Multi-Token Prediction. The drafter proposes K-1 tokens per round; the target verifies all in one parallel forward and accepts only those matching its own argmax. Output is bit-exact identical to standard greedy generation.

# Inference with pretrained drafter (~35% acceptance on generic prompts)
gemma4-cli generate \
  --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-bf16 \
  --draft-model google/gemma-4-E2B-it-assistant \
  --temperature 0 \
  "Your prompt"

# Chat mode (multi-turn)
gemma4-cli chat \
  --model-path ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-bf16 \
  --draft-model google/gemma-4-E2B-it-assistant

The pretrained drafter is generic — to actually win throughput, fine-tune it on the kind of text your target produces. Self-distillation against argmax(target_logits):

# Train (11 min for 2000 iter on 4.3k samples, batch=4)
gemma4-cli mtp-train \
  --target ~/Library/Caches/models/mlx-community/gemma-4-e2b-it-bf16 \
  --drafter google/gemma-4-E2B-it-assistant \
  --data my_corpus/train.jsonl \
  --valid-data my_corpus/valid.jsonl \
  --output ./my-drafter \
  --iterations 2000 --batch-size 4 \
  --steps-per-valid 200

# Inference with fine-tuned drafter + optional LoRA on target
gemma4-cli mtp-generate \
  --target mlx-community/gemma-4-e2b-it-bf16 \
  --drafter google/gemma-4-E2B-it-assistant \
  --drafter-path ./my-drafter/drafter.best.safetensors \
  --full-lm-head \
  --adapter-path ./my-target-lora \
  --prompt "Your domain-specific query"

Dataset format: same JSONL conventions as LoRA training ({"text": "..."} or {"messages": [{"role": ..., "content": ...}]}).

Greedy equivalence preserved (output bit-exact identical). See PR #25 for full bench.

gemma4-cli mtp-smoke <repo>              # validate drafter weights load cleanly
gemma4-cli mtp-forward                   # 1-round drafter parity test
gemma4-cli mtp-generate --compare        # bit-exact equivalence vs standard
gemma4-cli mtp-diag-verify               # sequential vs parallel hidden diff (advanced)

import Gemma4Swift

// Load a model (handles registration + tokenizer automatically)
let pipeline = Gemma4Pipeline()
try await pipeline.load(.e2b4bit)

// Or download + load in one call (no MLXLMCommon import needed)
try await pipeline.load(.e2b4bit, downloadIfNeeded: true) { progress in
    print("Downloading: \(Int(progress.fraction * 100))% — \(progress.currentFile)")
}

// Or from a custom local path
try await pipeline.load(from: URL(fileURLWithPath: "/path/to/model"))

// Chat
let response = try await pipeline.chat(prompt: "Hello!")

// Streaming
let stream = try pipeline.chatStream(prompt: "Write a poem")
for try await token in stream {
    print(token, terminator: "")
}

// Multi-turn
let followUp = try await pipeline.continueChat(prompt: "Make it shorter")

No need to import MLXLMCommon — Gemma4Pipeline.load() handles registration, tokenizer loading, and model container setup internally.

Gemma 4 models may spontaneously generate thinking blocks (<|channel>thought...). Use Gemma4TokenFilter to control this:

import Gemma4Swift

// Filter thinking (default) — only response content is emitted
let filter = Gemma4TokenFilter(mode: .disabled)

// Pass-through — raw tokens including thinking blocks
let filter = Gemma4TokenFilter(mode: .enabled)

// Structured — separate thinking and response
let filter = Gemma4TokenFilter(mode: .structured)

// In your generation loop:
let output = filter.process(tokenId: tokenId, text: decodedText)
if !output.isEmpty { print(output, terminator: "") }

// After generation (structured mode):
let result = filter.structuredResponse()
print("Thinking: \(result.thinking ?? "none")")
print("Response: \(result.response)")

Gemma4Swift/
├── Configuration/       # Model configs (text, vision, audio)
├── TextModel/           # Decoder layers, attention, MLP, MoE, per-layer inputs
├── RoPE/                # ProportionalRoPE with partial rotation
├── VisionEncoder/       # SigLIP: patch embed, 2D RoPE, pooler, head_dim padding
├── AudioEncoder/        # Conformer: SubSampleConv, chunked attention, rel positions
├── VideoProcessor/      # AVAsset frame extraction, ~1fps, aspect-ratio resize
├── Multimodal/          # Embedding fusion via masked_scatter
├── LoRA/                # LoRA/DoRA/Full SFT training, adapter load/fuse/unload
├── TurboQuant/          # MSE codec, Metal kernels, chunked prefill attention
├── Pipeline/            # High-level API, processors, token filter, registration
├── Norms/               # RMSNormNoScale, RMSNormZeroShift
└── Utils/               # Weight sanitizer, profiling toolkit

The project includes a custom TurboQuant implementation (rotation + Beta-optimal codebook) in TurboQuant/. After thorough audit, we recommend migrating to mlx-swift-lm's native QuantizedKVCache instead:

// Native KV cache quantization — no custom code needed
let params = GenerateParameters(kvBits: 4, kvGroupSize: 64, quantizedKVStart: 5000)

Why migrate: TurboQuant's theoretical 3.85x compression is real, but runtime intermediate tensor materialization erases memory gains. The disabled Fast Hadamard Transform (O(D^2) dense rotation instead of O(D log D)) adds significant compute overhead. Meanwhile, mlx-swift-lm ships battle-tested 4/8-bit KV quantization with Metal-accelerated quantizedMM(), automatic attention routing, and zero maintenance.

The TurboQuant/ directory is retained for reference but is not used in the default inference pipeline.

• Gemma 4 ≠ Gemma 3n: No AltUp, no Laurel blocks, no activation sparsity. Simpler decoder with global_head_dim, partial_rotary_factor, use_double_wide_mlp, and optional K=V attention.
• Registration-based: Registers "gemma4" and "gemma4_text" model types into mlx-swift-lm's LLMTypeRegistry.
• Multimodal via masked_scatter: Image/audio/video embeddings replace special token positions in the text embedding sequence.
• Video aligned with Google reference: video_token (258884), 70 soft tokens/frame, MM:SS timestamps, ~1fps sampling.
• Audio aligned with Google reference: Conformer with relative position embeddings, causal chunked attention, mel spectrogram matching feature_extraction_gemma4.py.
• Thinking mode handling: Gemma4TokenFilter detects <|channel>thought/<|channel>response blocks and filters or separates them at the API level.
• ProportionalRoPE: Only 25% of head dimensions get rotary encoding for full-attention layers.
• Vision head_dim padding: Pads non-standard head dimensions (72 → 80) for fused SDPA to avoid NaN from all-masked padding rows.
• No HuggingFace SDK dependency: Direct HTTPS downloads + local tokenizer loading.

• Google Gemma 4 — Original model architecture
• mlx-swift — Apple MLX framework for Swift
• mlx-swift-lm — LLM infrastructure
• mlx-vlm — Python reference implementation
• swift-transformers — Tokenizer support

MIT License — See LICENSE file.