dflash_distill_mlx.py

#!/usr/bin/env python3
"""Distill a classic Qwen3.5 DFlash draft against Ornith with MLX.

The pipeline creates target-generated responses, caches frozen Ornith hidden
features, then fine-tunes the draft with the block objective from the DFlash
paper. Target embeddings and the target LM head remain frozen.
"""

import argparse
import json
import math
import random
import shutil
import time
from pathlib import Path

import mlx.core as mx
from mlx import nn
import mlx.optimizers as optim
from mlx.utils import tree_flatten
from datasets import load_dataset
from huggingface_hub import snapshot_download
from mlx_lm import generate, load as load_target
from mlx_lm.models.cache import make_prompt_cache
from mlx_lm.sample_utils import make_sampler

from dflash.model_mlx import _patch_model, load_draft


def parse_args():
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="command", required=True)

    prepare = sub.add_parser("prepare")
    prepare.add_argument("--target", default="ornith-ai/Ornith-1.5-9B-MLX-4bit")
    prepare.add_argument("--draft", default="z-lab/Qwen3.5-9B-DFlash")
    prepare.add_argument("--output", type=Path, required=True)
    prepare.add_argument("--samples", type=int, default=128)
    prepare.add_argument("--eval-samples", type=int, default=16)
    prepare.add_argument("--max-new-tokens", type=int, default=192)
    prepare.add_argument("--max-sequence-tokens", type=int, default=384)
    prepare.add_argument("--dataset", default="tatsu-lab/alpaca")
    prepare.add_argument("--seed", type=int, default=20260821)
    recache = sub.add_parser("recache")
    recache.add_argument("--target", default="ornith-ai/Ornith-1.5-9B")
    recache.add_argument("--draft", default="z-lab/Qwen3.5-9B-DFlash")
    recache.add_argument("--data", type=Path, required=True)


    train = sub.add_parser("train")
    train.add_argument("--target", default="ornith-ai/Ornith-1.5-9B-MLX-4bit")
    train.add_argument("--draft", default="z-lab/Qwen3.5-9B-DFlash")
    train.add_argument("--data", type=Path, required=True)
    train.add_argument("--output", type=Path, required=True)
    train.add_argument("--steps", type=int, default=768)
    train.add_argument("--block-size", type=int, default=8)
    train.add_argument("--learning-rate", type=float, default=2e-5)
    train.add_argument("--weight-decay", type=float, default=0.01)
    train.add_argument("--warmup-ratio", type=float, default=0.04)
    train.add_argument("--clip-grad", type=float, default=1.0)
    train.add_argument("--loss-gamma", type=float, default=4.0)
    train.add_argument("--eval-every", type=int, default=96)
    train.add_argument("--save-every", type=int, default=192)
    train.add_argument(
        "--train-scope", choices=("all", "projection"), default="all"
    )
    train.add_argument("--seed", type=int, default=20260821)

    evaluate = sub.add_parser("evaluate")
    evaluate.add_argument("--target", default="ornith-ai/Ornith-1.5-9B-MLX-4bit")
    evaluate.add_argument("--draft", required=True)
    evaluate.add_argument("--data", type=Path, required=True)
    evaluate.add_argument("--block-size", type=int, default=8)
    evaluate.add_argument("--anchors", type=int, default=64)
    evaluate.add_argument("--seed", type=int, default=20260821)
    return parser.parse_args()


def dataset_prompts(name, limit):
    dataset = load_dataset(name, split="train", streaming=True)
    seen = set()
    for row in dataset:
        instruction = str(row.get("instruction") or row.get("prompt") or "").strip()
        extra = str(row.get("input") or "").strip()
        if not instruction or instruction in seen:
            continue
        seen.add(instruction)
        if extra:
            instruction = f"{instruction}\n\nInput:\n{extra}"
        yield instruction
        if len(seen) >= limit:
            return


def token_list(value):
    if hasattr(value, "tolist"):
        value = value.tolist()
    if value and isinstance(value[0], list):
        value = value[0]
    return [int(token) for token in value]


def target_features(model, tokens, layer_ids):
    cache = make_prompt_cache(model)
    model(mx.array(tokens, dtype=mx.int32)[None], cache)
    hidden = mx.concatenate(model._hidden_states, axis=-1)[0]
    hidden = mx.stop_gradient(hidden.astype(mx.bfloat16))
    mx.eval(hidden)
    return hidden


def prepare_data(args):
    random.seed(args.seed)
    mx.random.seed(args.seed)
    args.output.mkdir(parents=True, exist_ok=True)
    samples_dir = args.output / "samples"
    samples_dir.mkdir(exist_ok=True)

    model, tokenizer = load_target(args.target)
    draft = load_draft(args.draft)
    _patch_model(model, draft.config.target_layer_ids)
    model.eval()
    sampler = make_sampler(temp=0.6, top_p=0.95, top_k=20)
    total = args.samples + args.eval_samples
    manifest = []
    started = time.perf_counter()

    for index, instruction in enumerate(dataset_prompts(args.dataset, total)):
        messages = [{"role": "user", "content": instruction}]
        prompt_ids = token_list(tokenizer.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True
        ))
        response = generate(
            model,
            tokenizer,
            prompt=prompt_ids,
            max_tokens=args.max_new_tokens,
            sampler=sampler,
            verbose=False,
        )
        response_ids = token_list(tokenizer.encode(response, add_special_tokens=False))
        tokens = (prompt_ids + response_ids)[:args.max_sequence_tokens]
        if len(tokens) < len(prompt_ids) + 9:
            continue
        hidden = target_features(model, tokens, draft.config.target_layer_ids)
        sample_name = f"{len(manifest):05d}.safetensors"
        mx.save_safetensors(str(samples_dir / sample_name), {
            "tokens": mx.array(tokens, dtype=mx.int32),
            "hidden": hidden,
            "prompt_length": mx.array([len(prompt_ids)], dtype=mx.int32),
        })
        split = "eval" if len(manifest) < args.eval_samples else "train"
        manifest.append({
            "file": sample_name,
            "split": split,
            "instruction": instruction,
            "response": response,
            "tokens": len(tokens),
            "prompt_tokens": len(prompt_ids),
        })
        elapsed = time.perf_counter() - started
        print(f"[{len(manifest):3d}/{total}] {split:5s} tokens={len(tokens):3d} "
              f"elapsed={elapsed / 60:.1f}m", flush=True)
        mx.clear_cache()
        if len(manifest) >= total:
            break

    if len(manifest) < total:
        raise RuntimeError(f"prepared {len(manifest)} samples, expected {total}")
    with (args.output / "manifest.jsonl").open("w", encoding="utf-8") as out:
        for row in manifest:
            out.write(json.dumps(row, ensure_ascii=False) + "\n")
    metadata = {
        "target": args.target,
        "draft": args.draft,
        "dataset": args.dataset,
        "seed": args.seed,
        "samples": args.samples,
        "eval_samples": args.eval_samples,
    }
    (args.output / "config.json").write_text(
        json.dumps(metadata, indent=2) + "\n"
    )

def recache_features(args):
    model, _ = load_target(args.target)
    draft = load_draft(args.draft)
    _patch_model(model, draft.config.target_layer_ids)
    model.eval()
    rows = read_manifest(args.data, "eval") + read_manifest(args.data, "train")
    started = time.perf_counter()
    for index, row in enumerate(rows, 1):
        path = args.data / "samples" / row["file"]
        sample = mx.load(str(path))
        tokens = mx.array(sample["tokens"])
        prompt_length = mx.array(sample["prompt_length"])
        mx.eval(tokens, prompt_length)
        hidden = target_features(
            model, tokens.tolist(), draft.config.target_layer_ids
        )
        temporary = path.with_suffix(".tmp.safetensors")
        mx.save_safetensors(str(temporary), {
            "tokens": tokens,
            "hidden": hidden,
            "prompt_length": prompt_length,
        })
        temporary.replace(path)
        print(f"[{index:3d}/{len(rows)}] {row['file']} "
              f"elapsed={(time.perf_counter() - started) / 60:.1f}m",
              flush=True)
        mx.clear_cache()



def read_manifest(data_dir, split):
    rows = []
    with (data_dir / "manifest.jsonl").open(encoding="utf-8") as src:
        for line in src:
            row = json.loads(line)
            if row["split"] == split:
                rows.append(row)
    if not rows:
        raise RuntimeError(f"no {split} samples in {data_dir}")
    return rows


def load_sample(data_dir, row):
    return mx.load(str(data_dir / "samples" / row["file"]))


def choose_anchor(sample, block_size, rng):
    tokens = sample["tokens"]
    prompt_length = int(sample["prompt_length"][0].item())
    last = int(tokens.shape[0]) - block_size
    if last < prompt_length:
        raise RuntimeError("sample response is shorter than the draft block")
    return rng.randint(prompt_length, last)


def bind_and_freeze(draft, target, scope="all"):
    draft.bind(target)
    draft.freeze()
    draft.fc.unfreeze()
    draft.hidden_norm.unfreeze()
    if scope == "all":
        for layer in draft.layers:
            layer.unfreeze()
        draft.norm.unfreeze()
    draft.train()
    target.eval()


def loss_fn(draft, sample, anchor, block_size, loss_gamma):
    tokens = sample["tokens"]
    target_hidden = sample["hidden"][None, :anchor]
    draft_input = mx.concatenate((
        tokens[anchor:anchor + 1],
        mx.full((block_size - 1,), draft.config.mask_token_id, dtype=mx.int32),
    ))[None]
    labels = tokens[anchor + 1:anchor + block_size][None]
    logits = draft(draft_input, target_hidden, draft.make_cache(), logits_start=1)
    losses = nn.losses.cross_entropy(logits, labels, reduction="none")
    positions = mx.arange(block_size - 1, dtype=mx.float32)
    weights = mx.exp(-positions / loss_gamma)
    return mx.sum(losses * weights[None]) / mx.sum(weights)


def acceptance_for_sample(draft, sample, anchor, block_size):
    tokens = sample["tokens"]
    target_hidden = sample["hidden"][None, :anchor]
    draft_input = mx.concatenate((
        tokens[anchor:anchor + 1],
        mx.full((block_size - 1,), draft.config.mask_token_id, dtype=mx.int32),
    ))[None]
    logits = draft(draft_input, target_hidden, draft.make_cache(), logits_start=1)
    predictions = mx.argmax(logits, axis=-1)[0]
    expected = tokens[anchor + 1:anchor + block_size]
    matches = mx.equal(predictions, expected).tolist()
    accepted = 0
    for match in matches:
        if not match:
            break
        accepted += 1
    return accepted


def evaluate_draft(draft, data_dir, rows, block_size, anchors, seed):
    rng = random.Random(seed)
    draft.eval()
    accepted = []
    exact = 0
    for index in range(anchors):
        row = rows[index % len(rows)]
        sample = load_sample(data_dir, row)
        anchor = choose_anchor(sample, block_size, rng)
        count = acceptance_for_sample(draft, sample, anchor, block_size)
        accepted.append(count)
        exact += count == block_size - 1
        if index % 8 == 7:
            mx.clear_cache()
    draft.train()
    mean = sum(accepted) / len(accepted)
    return {
        "anchors": len(accepted),
        "mean_accepted_drafts": mean,
        "mean_cycle_tokens": mean + 1.0,
        "full_blocks": exact,
        "full_block_rate": exact / len(accepted),
    }


def draft_source_path(draft_id):
    return Path(snapshot_download(draft_id, allow_patterns=["*.safetensors", "*.json"]))


def save_draft(draft, draft_id, output, metrics, training):
    output.mkdir(parents=True, exist_ok=True)
    source = draft_source_path(draft_id)
    original = {
        key
        for file in source.glob("*.safetensors")
        for key in mx.load(str(file)).keys()
    }
    parameters = dict(tree_flatten(draft.parameters()))
    missing = sorted(original - parameters.keys())
    if missing:
        raise RuntimeError(f"trained model is missing original weights: {missing}")
    weights = {key: parameters[key] for key in sorted(original)}
    mx.eval(weights)
    mx.save_safetensors(str(output / "model.safetensors"), weights)
    shutil.copy2(source / "config.json", output / "config.json")
    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
    (output / "training.json").write_text(json.dumps(training, indent=2) + "\n")


def learning_rate(step, total, peak, warmup_ratio):
    warmup = max(1, round(total * warmup_ratio))
    if step < warmup:
        return peak * (step + 1) / warmup
    progress = (step - warmup) / max(1, total - warmup)
    return peak * 0.5 * (1.0 + math.cos(math.pi * progress))


def train_draft(args):
    rng = random.Random(args.seed)
    mx.random.seed(args.seed)
    train_rows = read_manifest(args.data, "train")
    eval_rows = read_manifest(args.data, "eval")
    target, _ = load_target(args.target)
    draft = load_draft(args.draft)
    bind_and_freeze(draft, target, args.train_scope)

    optimizer = optim.AdamW(
        learning_rate=args.learning_rate,
        weight_decay=args.weight_decay,
    )
    loss_and_grad = nn.value_and_grad(draft, loss_fn)
    baseline = evaluate_draft(
        draft, args.data, eval_rows, args.block_size,
        min(64, len(eval_rows) * 4), args.seed,
    )
    print(f"baseline {json.dumps(baseline, sort_keys=True)}", flush=True)
    best = baseline
    started = time.perf_counter()

    for step in range(args.steps):
        row = train_rows[step % len(train_rows)]
        if step and step % len(train_rows) == 0:
            rng.shuffle(train_rows)
        sample = load_sample(args.data, row)
        anchor = choose_anchor(sample, args.block_size, rng)
        rate = learning_rate(step, args.steps, args.learning_rate, args.warmup_ratio)
        optimizer.learning_rate = mx.array(rate)
        loss, grads = loss_and_grad(
            draft, sample, anchor, args.block_size, args.loss_gamma
        )
        grads, grad_norm = optim.clip_grad_norm(grads, args.clip_grad)
        optimizer.update(draft, grads)
        mx.eval(loss, grad_norm, draft.parameters(), optimizer.state)

        if step == 0 or (step + 1) % 8 == 0:
            elapsed = time.perf_counter() - started
            print(f"step={step + 1}/{args.steps} loss={loss.item():.5f} "
                  f"grad={grad_norm.item():.3f} lr={rate:.3e} "
                  f"steps_s={(step + 1) / elapsed:.3f}", flush=True)
        if (step + 1) % args.eval_every == 0 or step + 1 == args.steps:
            metrics = evaluate_draft(
                draft, args.data, eval_rows, args.block_size,
                min(64, len(eval_rows) * 4), args.seed,
            )
            print(f"eval step={step + 1} {json.dumps(metrics, sort_keys=True)}", flush=True)
            if metrics["mean_accepted_drafts"] > best["mean_accepted_drafts"]:
                best = metrics
                save_draft(draft, args.draft, args.output / "best", metrics, json_safe_args(args))
        if (step + 1) % args.save_every == 0:
            save_draft(
                draft, args.draft, args.output / f"step-{step + 1:06d}",
                metrics if "metrics" in locals() else baseline, json_safe_args(args),
            )
        mx.clear_cache()

    final_metrics = evaluate_draft(
        draft, args.data, eval_rows, args.block_size,
        min(64, len(eval_rows) * 4), args.seed,
    )
    save_draft(draft, args.draft, args.output / "final", final_metrics, json_safe_args(args))
    summary = {"baseline": baseline, "best": best, "final": final_metrics}
    (args.output / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    print(json.dumps(summary, indent=2), flush=True)


def evaluate_command(args):
    rows = read_manifest(args.data, "eval")
    target, _ = load_target(args.target)
    draft = load_draft(args.draft)
    bind_and_freeze(draft, target)
    metrics = evaluate_draft(
        draft, args.data, rows, args.block_size, args.anchors, args.seed
    )
    print(json.dumps(metrics, indent=2))


def json_safe_args(args):
    values = vars(args).copy()
    for key, value in values.items():
        if isinstance(value, Path):
            values[key] = str(value)
    return values


def main():
    args = parse_args()
    if args.command == "prepare":
        prepare_data(args)
    elif args.command == "train":
        train_draft(args)
    elif args.command == "recache":
        recache_features(args)
    else:
        evaluate_command(args)


if __name__ == "__main__":
    main()