code/diffusion_lora_train.py

#!/usr/bin/env python3
"""
DiffusionGemma block-diffusion QLoRA trainer for Apple Silicon (MLX).

There is no diffusion-aware trainer in the MLX ecosystem (mlx-vlm's SFT trainer
computes autoregressive next-token loss — the wrong objective for this model).
This implements the ground-truth DiffusionGemma SFT recipe, established from
Google's JAX reference (gemma hackable_diffusion_adapter), NeMo Automodel's
diffusion_gemma_lora.yaml, and Unsloth's notebook:

  * Corruption: D3PM-uniform. One t ~ U(eps, 1) per example (eps=0.001);
    each canvas position independently replaced with prob t by a token drawn
    uniformly from [0, vocab). NO mask token — matches inference, which
    initializes/renoises the canvas with mx.random.randint over the vocab.
  * Canvas: response-relative 256-token grid. Response tokens (ending with
    <turn|> id 106) are EOS-filled (id 1) to the canvas boundary; the fill is
    supervised and attended (the model learns termination).
  * Loss: flat UNWEIGHTED cross-entropy over ALL canvas positions — corrupted
    and uncorrupted alike (Google NoWeightDiscreteLoss; the corrupted-only
    variant was a documented NeMo bug). No 1/t weighting (that is the
    absorbing-kernel ELBO weight; does not apply to the uniform kernel).
  * v1 simplifications (Unsloth-proven, 1.5%->89.5% on Sudoku): no co-trained
    encoder AR loss, no self-conditioning passes. v2 candidates documented.
  * LoRA (NeMo parity): r=16 alpha=32 on self_attn {q,k,v,o}_proj (v_proj only
    exists on the 25 sliding layers) + dense mlp {gate,up,down}_proj.
    MoE experts, router, embeddings frozen. QLoRA: base stays 4-bit quantized
    (mlx_lm LoRALinear.from_base is QuantizedLinear-aware).
  * Optimizer (NeMo): AdamW (bias-corrected) lr 1.5e-4, betas (0.95, 0.99),
    eps 1e-8, wd 1e-4, grad-norm clip 1.0, 25-step warmup, cosine to 1.5e-5,
    grad-accum 8 at micro-batch 1. Step count via --steps (the cosine horizon
    scales with it). Note: only LoRA weights are checkpointed — crash-resume
    restarts Adam moments (recovers over ~100 steps); resumed runs are not
    sample-equivalent to uninterrupted ones.

Adapters are saved in mlx-vlm-compatible format (adapters.safetensors with
full dotted root paths + adapter_config.json carrying lora_parameters.keys),
so mlx_vlm.trainer.utils.apply_lora_layers(model, adapter_dir) can reload them.

Usage:
  python3 diffusion_lora_train.py \
      --model ./diffusiongemma-26B-A4B-it-4bit \
      --data ./data \
      --adapter-path ./adapters/my-adapter \
      --smoke              # 3 fwd/bwd iters + sanity prints, then exit
"""

import argparse
import json
import time
from pathlib import Path

import numpy as np

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from mlx.utils import tree_flatten, tree_map

from mlx_lm.tuner.lora import LoRALinear
from mlx_vlm.utils import load as vlm_load

VOCAB_SIZE = 262144
CANVAS_LEN = 256
BOS_ID = 2
EOS_FILL_ID = 1     # <eos> — Google's CanvasChunker EOS-fill token
TURN_END_ID = 106   # <turn|> — response terminator from the chat template
SOFTCAP = 30.0
T_EPS = 0.001       # Google SafeSpan eps 1e-4, NeMo 0.001; we use NeMo's


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True)
    p.add_argument("--data", required=True, help="dir with train/valid jsonl of {prompt, response}")
    p.add_argument("--adapter-path", required=True)
    p.add_argument("--steps", type=int, default=800)
    p.add_argument("--accum", type=int, default=8)
    p.add_argument("--lr", type=float, default=1.5e-4)
    p.add_argument("--min-lr", type=float, default=1.5e-5)
    p.add_argument("--warmup", type=int, default=25)
    p.add_argument("--rank", type=int, default=16)
    p.add_argument("--alpha", type=float, default=32.0)
    p.add_argument("--max-prompt-tokens", type=int, default=1920)
    p.add_argument("--val-every", type=int, default=50)
    p.add_argument("--val-samples", type=int, default=32)
    p.add_argument("--save-every", type=int, default=100)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--grad-checkpoint", action="store_true")
    p.add_argument("--cache-limit-gb", type=float, default=2.0)
    p.add_argument("--smoke", action="store_true", help="3 iters of sanity checks, then exit")
    p.add_argument("--resume-file", default=None, help="adapter safetensors to resume from")
    p.add_argument("--start-step", type=int, default=0, help="optimizer step to resume at (schedule offset)")
    return p.parse_args()


# ── LoRA wiring ──────────────────────────────────────────────────────────────

ATTN_PROJS = ("q_proj", "k_proj", "v_proj", "o_proj")
MLP_PROJS = ("gate_proj", "up_proj", "down_proj")


def apply_lora(model, rank, alpha):
    """Freeze everything, wrap attention + dense-MLP projections with LoRA.

    Stock mlx_vlm get_peft_model crashes on diffusion_gemma (_LanguageModelView
    is not an nn.Module), so we wire LoRALinear.from_base manually.
    """
    model.freeze()
    scale = alpha / rank
    keys = []
    for i, layer in enumerate(model.model.decoder.layers):
        for proj in ATTN_PROJS:
            lin = getattr(layer.self_attn, proj)
            if lin is None:  # full-attention layers have v_proj=None (values reuse keys)
                continue
            setattr(layer.self_attn, proj, LoRALinear.from_base(lin, r=rank, scale=scale))
            keys.append(f"model.decoder.layers.{i}.self_attn.{proj}")
        for proj in MLP_PROJS:
            lin = getattr(layer.mlp, proj)
            setattr(layer.mlp, proj, LoRALinear.from_base(lin, r=rank, scale=scale))
            keys.append(f"model.decoder.layers.{i}.mlp.{proj}")
    return keys, scale


def save_adapter(model, adapter_dir, keys, rank, scale, step, extra=None):
    adapter_dir = Path(adapter_dir)
    adapter_dir.mkdir(parents=True, exist_ok=True)
    weights = dict(tree_flatten(model.trainable_parameters()))
    name = "adapters.safetensors" if step is None else f"{step:07d}_adapters.safetensors"
    mx.save_safetensors(str(adapter_dir / name), weights)
    config = {
        "fine_tune_type": "lora",
        "model_type": "diffusion_gemma",
        "num_layers": -1,
        "lora_parameters": {"rank": rank, "scale": scale, "dropout": 0.0, "keys": keys},
        "training": extra or {},
    }
    (adapter_dir / "adapter_config.json").write_text(json.dumps(config, indent=2))


# ── Data ─────────────────────────────────────────────────────────────────────

def load_split(path, tokenizer, max_prompt_tokens):
    samples, skipped = [], 0
    with open(path) as f:
        for line in f:
            obj = json.loads(line)
            p_ids = [BOS_ID] + tokenizer.encode(obj["prompt"], add_special_tokens=False)
            r_ids = tokenizer.encode(obj["response"], add_special_tokens=False)
            if len(p_ids) > max_prompt_tokens or len(r_ids) > CANVAS_LEN:
                skipped += 1
                continue
            clean = r_ids + [EOS_FILL_ID] * (CANVAS_LEN - len(r_ids))
            samples.append((np.asarray(p_ids, np.int32), np.asarray(clean, np.int32), len(r_ids)))
    return samples, skipped


def corrupt(clean, rng):
    """D3PM-uniform: t ~ U(eps,1); replace each position w.p. t by uniform token."""
    t = T_EPS + (1.0 - T_EPS) * rng.random()
    mask = rng.random(clean.shape[0]) < t
    noise = rng.integers(0, VOCAB_SIZE, clean.shape[0], dtype=np.int32)
    return np.where(mask, noise, clean).astype(np.int32), t


# ── Loss ─────────────────────────────────────────────────────────────────────

def make_loss_fn():
    def loss_fn(model, prompt_ids, canvas_ids, targets):
        hidden, _ = model.model(input_ids=prompt_ids, canvas_ids=canvas_ids)
        logits = model.model.decoder.embed_tokens.as_linear(hidden)
        logits = mx.tanh(logits.astype(mx.float32) / SOFTCAP) * SOFTCAP
        ce = nn.losses.cross_entropy(logits, targets, reduction="none")
        return ce.mean()  # all canvas positions supervised, unweighted
    return loss_fn


def main():
    args = parse_args()
    mx.random.seed(args.seed)
    rng = np.random.default_rng(args.seed)
    mx.set_cache_limit(int(args.cache_limit_gb * 1024**3))

    print(f"[load] {args.model}", flush=True)
    model, processor = vlm_load(args.model)
    tokenizer = processor.tokenizer

    # the corruption/canvas constants are checkpoint-specific — fail loudly on a mismatch
    text_cfg = getattr(model.config, "text_config", model.config)
    cfg_vocab = getattr(text_cfg, "vocab_size", VOCAB_SIZE)
    cfg_canvas = getattr(model.config, "canvas_length", CANVAS_LEN)
    assert cfg_vocab == VOCAB_SIZE, f"checkpoint vocab_size {cfg_vocab} != {VOCAB_SIZE}"
    assert cfg_canvas == CANVAS_LEN, f"checkpoint canvas_length {cfg_canvas} != {CANVAS_LEN}"

    keys, scale = apply_lora(model, args.rank, args.alpha)
    n_train = sum(v.size for _, v in tree_flatten(model.trainable_parameters()))
    print(f"[lora] {len(keys)} layers wrapped, {n_train/1e6:.2f}M trainable params, scale={scale}", flush=True)

    if args.resume_file:
        ckpt_keys = set(mx.load(args.resume_file).keys())
        trainable_keys = set(dict(tree_flatten(model.trainable_parameters())).keys())
        missing = trainable_keys - ckpt_keys
        assert not missing, f"resume checkpoint missing {len(missing)} trainable keys, e.g. {sorted(missing)[:3]}"
        model.load_weights(args.resume_file, strict=False)
        print(f"[resume] loaded {args.resume_file} ({len(ckpt_keys)} keys), continuing at step {args.start_step}", flush=True)

    if args.grad_checkpoint:
        from mlx_vlm.trainer.utils import grad_checkpoint
        grad_checkpoint(model.model.decoder.layers[0])
        print("[mem] gradient checkpointing ON (class-wide DecoderLayer patch)", flush=True)

    model.train()

    data_dir = Path(args.data)
    train, n_skip_t = load_split(data_dir / "train.jsonl", tokenizer, args.max_prompt_tokens)
    valid, n_skip_v = load_split(data_dir / "valid.jsonl", tokenizer, args.max_prompt_tokens)
    print(f"[data] train={len(train)} (skipped {n_skip_t}), valid={len(valid)} (skipped {n_skip_v})", flush=True)

    # sanity: responses must end with <turn|> before the EOS fill
    _, clean0, rlen0 = train[0]
    assert clean0[rlen0 - 1] == TURN_END_ID, f"response does not end with <turn|>: {clean0[max(0,rlen0-5):rlen0]}"
    if rlen0 < CANVAS_LEN:
        assert clean0[rlen0] == EOS_FILL_ID

    # fixed validation corruptions for comparable val loss across steps
    vrng = np.random.default_rng(args.seed + 1)
    val_idx = vrng.choice(len(valid), size=min(args.val_samples, len(valid)), replace=False)
    val_set = []
    for i in val_idx:
        p, clean, _ = valid[i]
        canvas, _t = corrupt(clean, vrng)
        val_set.append((p, canvas, clean))

    loss_fn = make_loss_fn()
    loss_and_grad = nn.value_and_grad(model, loss_fn)

    # linear covers counters 0..warmup-1 (ends at lr*(w-1)/w); cosine takes over at
    # counter==warmup starting exactly at lr — a continuous ramp with increment lr/w
    base_schedule = optim.join_schedules(
        [optim.linear_schedule(0.0, args.lr, args.warmup),
         optim.cosine_decay(args.lr, max(1, args.steps - args.warmup), args.min_lr)],
        [args.warmup],
    )
    # on resume the optimizer's internal counter restarts at 0 — offset the schedule
    # so the LR continues from where the crashed run left off (Adam moments are lost)
    schedule = (lambda s: base_schedule(s + args.start_step)) if args.start_step else base_schedule
    # bias_correction=True for parity with torch AdamW (NeMo recipe): mlx defaults
    # to False, which inflates the effective step ~1.3-1.5x during early-mid training
    try:
        optimizer = optim.AdamW(learning_rate=schedule, betas=[0.95, 0.99], eps=1e-8,
                                weight_decay=1e-4, bias_correction=True)
    except TypeError:  # older mlx without the kwarg
        optimizer = optim.AdamW(learning_rate=schedule, betas=[0.95, 0.99], eps=1e-8, weight_decay=1e-4)
        print("[warn] mlx AdamW lacks bias_correction — effective LR inflated early in the run", flush=True)

    def val_loss():
        model.eval()
        tot = 0.0
        for p, canvas, clean in val_set:
            l = loss_fn(model, mx.array(p[None]), mx.array(canvas[None]), mx.array(clean[None]))
            mx.eval(l)
            tot += l.item()
        model.train()
        return tot / len(val_set)

    if args.smoke:
        print("[smoke] 3 fwd/bwd iters", flush=True)
        for it in range(3):
            p, clean, rlen = train[it]
            canvas, t = corrupt(clean, rng)
            t0 = time.time()
            (loss), grads = loss_and_grad(model, mx.array(p[None]), mx.array(canvas[None]), mx.array(clean[None]))
            grads, gnorm = optim.clip_grad_norm(grads, 1.0)
            mx.eval(loss, grads)
            dt = time.time() - t0
            print(f"  it={it} t={t:.3f} resp_len={rlen} prompt_len={len(p)} "
                  f"loss={loss.item():.4f} gnorm={gnorm.item():.3f} {dt:.1f}s "
                  f"peak={mx.get_peak_memory()/1e9:.1f}GB", flush=True)
        print(f"[smoke] initial val loss = {val_loss():.4f}", flush=True)
        print("[smoke] OK", flush=True)
        return

    print(f"[train] steps={args.steps} accum={args.accum} lr={args.lr} start={args.start_step}", flush=True)
    log_path = Path(args.adapter_path) / "train_log.jsonl"
    Path(args.adapter_path).mkdir(parents=True, exist_ok=True)
    if args.start_step:  # fresh shuffle stream on resume; exact replay not required for SFT
        rng = np.random.default_rng(args.seed + args.start_step)
    order = rng.permutation(len(train))
    cursor = 0
    t_start = time.time()
    # carry best_val across crash-resumes so a post-resume regression can't
    # overwrite the published best adapter
    best_json = Path(args.adapter_path) / "best.json"
    best_val = float("inf")
    if args.start_step and best_json.exists():
        best_val = json.loads(best_json.read_text())["val_loss"]
        print(f"[resume] best_val carried over: {best_val:.4f}", flush=True)

    for step in range(args.start_step + 1, args.steps + 1):
        acc_grads = None
        acc_loss = 0.0
        for _ in range(args.accum):
            if cursor >= len(order):
                order = rng.permutation(len(train))
                cursor = 0
            p, clean, _ = train[order[cursor]]
            cursor += 1
            canvas, _t = corrupt(clean, rng)
            (loss), grads = loss_and_grad(model, mx.array(p[None]), mx.array(canvas[None]), mx.array(clean[None]))
            acc_grads = grads if acc_grads is None else tree_map(mx.add, acc_grads, grads)
            # materialize accumulated grads each micro-step: frees this step's
            # backward graph so memory stays at single-step peak (~28GB measured)
            # instead of stacking `accum` lazy backward graphs
            mx.eval(acc_grads, loss)
            acc_loss += loss.item()
        acc_grads = tree_map(lambda g: g / args.accum, acc_grads)
        acc_grads, gnorm = optim.clip_grad_norm(acc_grads, 1.0)
        optimizer.update(model, acc_grads)
        mx.eval(model.parameters(), optimizer.state)

        rec = {
            "step": step,
            "loss": acc_loss / args.accum,
            "gnorm": gnorm.item(),
            # optimizer's internal counter is 0-indexed AND starts at 0 on resume,
            # so the LR actually used at global step N is base_schedule(N-1-start_step+start_step)
            # = base_schedule applied to the optimizer counter offset by start_step
            "lr": float(base_schedule(mx.array(step - 1)).item()),
            "peak_gb": mx.get_peak_memory() / 1e9,
            "elapsed_s": round(time.time() - t_start, 1),
        }
        if step % args.val_every == 0 or step == args.steps:
            rec["val_loss"] = val_loss()
            if rec["val_loss"] < best_val:
                best_val = rec["val_loss"]
                save_adapter(model, args.adapter_path, keys, args.rank, scale, None,
                             extra={"step": step, "val_loss": rec["val_loss"], "best": True})
                best_json.write_text(json.dumps({"step": step, "val_loss": best_val}))
                rec["saved_best"] = True
        if step % args.save_every == 0:
            save_adapter(model, args.adapter_path, keys, args.rank, scale, step)
        with open(log_path, "a") as f:
            f.write(json.dumps(rec) + "\n")
        print(f"step {step}/{args.steps} loss={rec['loss']:.4f} "
              f"{'val=' + format(rec['val_loss'], '.4f') + ' ' if 'val_loss' in rec else ''}"
              f"gnorm={rec['gnorm']:.2f} peak={rec['peak_gb']:.1f}GB", flush=True)

    save_adapter(model, args.adapter_path, keys, args.rank, scale, args.steps,
                 extra={"step": args.steps, "final": True, "best_val": best_val})
    print(f"[done] best_val={best_val:.4f} adapter={args.adapter_path} "
          f"total={time.time()-t_start:.0f}s", flush=True)


if __name__ == "__main__":
    main()