# DiffusionGemma QLoRA on Apple Silicon — trainer + eval

The first known fine-tuning pipeline for [google/diffusiongemma-26B-A4B-it](https://huggingface.co/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)

```bash
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

```bash
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

```bash
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).