File size: 3,405 Bytes
60b165e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | # 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).
|