| """Matched full-pipeline optimization trials, with load and warm runs separate.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import json |
| import logging |
| from pathlib import Path |
|
|
| from h3_mlx.runtime import Pipeline |
| from h3_mlx.weights import Linear |
|
|
|
|
| class HalfProjection(Linear): |
| """Experimental FP16 GEMM, with BF16 rounding at every native boundary.""" |
|
|
| def __call__(self, x): |
| import mlx.core as mx |
|
|
| return super().__call__(x).astype(mx.bfloat16) |
|
|
|
|
| PROFILES = { |
| "quiet-baseline": {"wired_gb": 42, "cache_refined_text": False}, |
| "quiet-native": { |
| "wired_gb": 42, |
| "compile_blocks": True, |
| "eval_every": 0, |
| "prefetch_blocks": 4, |
| "vae_batch": 4, |
| "vae_compile": True, |
| }, |
| "quiet-packed": { |
| "wired_gb": 32, |
| "compile_blocks": True, |
| "eval_every": 0, |
| "prefetch_blocks": 4, |
| "weight_storage": "packed12", |
| "vae_batch": 4, |
| "vae_compile": True, |
| }, |
| "speed": { |
| "wired_gb": 42, |
| "compile_blocks": True, |
| "eval_every": 0, |
| "vae_batch": 4, |
| "vae_compile": True, |
| }, |
| "prefetch": { |
| "wired_gb": 42, |
| "compile_blocks": True, |
| "eval_every": 0, |
| "prefetch_blocks": 4, |
| "vae_batch": 4, |
| "vae_compile": True, |
| }, |
| "compact": { |
| "wired_gb": 32, |
| "compile_blocks": True, |
| "eval_every": 0, |
| "prefetch_blocks": 4, |
| "weight_storage": "packed12", |
| "vae_batch": 4, |
| "vae_compile": True, |
| }, |
| "baseline": {}, |
| "compiled": {"compile_blocks": True}, |
| "compiled-batched": { |
| "compile_blocks": True, |
| "eval_every": 0, |
| "keep_refiners": True, |
| }, |
| "compiled-fast": { |
| "compile_blocks": True, |
| "eval_every": 0, |
| "keep_refiners": True, |
| "vae_batch": 8, |
| }, |
| "fused-compiled": { |
| "compile_blocks": True, |
| "eval_every": 0, |
| "keep_refiners": True, |
| "vae_batch": 8, |
| "kernel": "metal-qk", |
| }, |
| "batched": {"eval_every": 0, "keep_refiners": True}, |
| "wired-baseline": {"wired_gb": 42}, |
| "wired-compiled": { |
| "wired_gb": 42, |
| "compile_blocks": True, |
| "eval_every": 0, |
| "keep_refiners": True, |
| }, |
| "wired-fast": { |
| "wired_gb": 42, |
| "compile_blocks": True, |
| "eval_every": 0, |
| "keep_refiners": True, |
| "vae_batch": 8, |
| }, |
| "packed12": { |
| "wired_gb": 34, |
| "compile_blocks": True, |
| "eval_every": 0, |
| "keep_refiners": True, |
| "weight_storage": "packed12", |
| }, |
| "packed12-fast": { |
| "wired_gb": 34, |
| "compile_blocks": True, |
| "eval_every": 0, |
| "keep_refiners": True, |
| "weight_storage": "packed12", |
| "vae_batch": 8, |
| }, |
| "five-pass": { |
| "wired_gb": 42, |
| "compile_blocks": True, |
| "eval_every": 0, |
| "keep_refiners": True, |
| "vae_batch": 8, |
| }, |
| "fp16": { |
| "allow_nonfinite": True, |
| "wired_gb": 42, |
| "compile_blocks": True, |
| "eval_every": 0, |
| "keep_refiners": True, |
| "vae_batch": 8, |
| }, |
| } |
|
|
|
|
| def main(a): |
| logging.basicConfig( |
| level=logging.INFO, format="%(asctime)s %(message)s", datefmt="%H:%M:%S" |
| ) |
| root = Path(a.output) |
| root.mkdir(parents=True, exist_ok=True) |
| cases = [ |
| c |
| for c in json.loads(Path("experiments/cases.json").read_text()) |
| if c["id"] in a.cases |
| ] |
| for name in a.profiles: |
| original_load = Linear.load.__func__ |
| if name == "fp16": |
|
|
| def half_load(cls, store, prefix, original_load=original_load, **kwargs): |
| import mlx.core as mx |
|
|
| layer = original_load(cls, store, prefix, **kwargs) |
| if prefix.startswith("blocks."): |
| layer = HalfProjection( |
| layer.weight.astype(mx.float16), |
| None if layer.bias is None else layer.bias.astype(mx.float16), |
| ) |
| return layer |
|
|
| Linear.load = classmethod(half_load) |
| pipeline = Pipeline(a.model, a.text_encoder, **PROFILES[name]) |
| try: |
| for case in cases: |
| for trial in range(a.repeat): |
| _, row = pipeline.generate( |
| case["prompt"], |
| width=case["width"], |
| height=case["height"], |
| seed=case["seed"], |
| images=case.get("images", ()), |
| save_trace=True, |
| steps=[0, 2, 3, 4, 5] if name == "five-pass" else None, |
| output=root / case["id"] / f"{name}-{trial}.png", |
| ) |
| row.update(profile=name, case=case["id"], trial=trial) |
| if name == "fp16": |
| row["experimental_precision"] = ( |
| "Main projection GEMMs FP16, native BF16 output rounding; norms, attention and other components retain the original policy" |
| ) |
| Path(row["output"]).with_suffix(".json").write_text( |
| json.dumps(row, indent=2) + "\n" |
| ) |
| with (root / "trials.jsonl").open("a") as f: |
| f.write(json.dumps(row) + "\n") |
| finally: |
| pipeline.clear() |
| del pipeline |
| gc.collect() |
| Linear.load = classmethod(original_load) |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--model", required=True) |
| p.add_argument("--text-encoder", required=True) |
| p.add_argument( |
| "--profiles", nargs="+", choices=PROFILES, default=["compiled-batched"] |
| ) |
| p.add_argument("--cases", nargs="+", default=["product"]) |
| p.add_argument("--repeat", type=int, default=3) |
| p.add_argument("--output", default="outputs/optimization-v2") |
| main(p.parse_args()) |
|
|
