DiffusionGemma QLoRA on Apple Silicon — trainer + eval

The first known fine-tuning pipeline for google/diffusiongemma-26B-A4B-it that runs entirely on Apple Silicon via MLX. No diffusion-aware trainer exists in the MLX ecosystem (mlx-lm has no support for the architecture; mlx-vlm is inference-only and its SFT trainer optimizes the wrong, autoregressive objective). These scripts implement the correct block-diffusion denoising objective.

Requirements

Apple Silicon Mac (tested: M2 Max, 64 GB). Training peaks ~17.5 GB with --grad-checkpoint.
pip install -U "mlx-vlm>=0.6.3" "mlx-lm>=0.31.3" mlx
The 4-bit base: mlx-community/diffusiongemma-26B-A4B-it-4bit (~15.7 GB).

Files

File	What it does
`diffusion_lora_train.py`	Block-diffusion QLoRA trainer (the core contribution)
`diffusion_eval.py`	Tool-selection benchmark for DiffusionGemma (diffusion generation)
`ar_eval.py`	Same benchmark for autoregressive baselines (mlx_lm) — fair cross-model comparison
`analyze_results.py`	Bootstrap CIs + paired significance + length stratification → report
`make_example_data.py`	Generates a synthetic toy dataset (the real corpus is private)
`run_chain.sh`	eval → train → eval orchestration with crash-resume

60-second smoke test (synthetic data)

python3 make_example_data.py --out ./data      # 120/24/24 synthetic examples
hf download mlx-community/diffusiongemma-26B-A4B-it-4bit --local-dir ./dg-4bit

# 3 forward/backward sanity iters, then exit
python3 diffusion_lora_train.py --model ./dg-4bit --data ./data \
    --adapter-path ./adapters/toy --smoke

Full training

python3 diffusion_lora_train.py --model ./dg-4bit --data ./data \
    --adapter-path ./adapters/toy --steps 250 --grad-checkpoint

Key recipe choices (verified against Google JAX / NVIDIA NeMo / Unsloth sources): D3PM-uniform corruption (random vocab tokens, not mask tokens), unweighted CE over the full 256-token canvas including supervised EOS-fill, LoRA r16/α32 on attention + dense-MLP (MoE experts/router frozen), bias-corrected AdamW 1.5e-4 → cosine.

Evaluation

python3 diffusion_eval.py --model ./dg-4bit --adapter ./adapters/toy \
    --test ./data/test.jsonl --out ./eval.json
python3 analyze_results.py --dir .          # report with CIs + significance

Operational gotchas on Apple Silicon (learned the hard way)

Train with --grad-checkpoint — faster here than the unchunked backward on this MoE, and ~17.5 GB peak.
macOS GPU watchdog kills long Metal command buffers (...ImpactingInteractivity). For training use MLX_MAX_OPS_PER_BUFFER=4 MLX_MAX_MB_PER_BUFFER=20 + the crash-resume loop in run_chain.sh.
For diffusion EVAL on long prompts use the Goldilocks MLX_MAX_OPS_PER_BUFFER=32 MLX_MAX_MB_PER_BUFFER=128: tiny buffers make long-context prefill crawl; no caps let it trip the watchdog. A per-sample --sample-timeout is a backstop.
Long-context limitation: DiffusionGemma is impractically slow to generate from ~900-token prompts on MLX today (the encoder prefill), even though training is fine — a real consideration for agentic use where prompts are long.

License

Apache-2.0 (matches the base model). Training corpus not included (private).