From 042a1d656384c6e88e145182cfc3d9a4403bafaa Mon Sep 17 00:00:00 2001 From: Rafal Date: Wed, 1 Jul 2026 09:26:01 +0000 Subject: [PATCH] feat(sa3): fp16 checkpoint converter (halves DiT on-disk size) The medium DiT ships as ~9 GB of float32 safetensors but DEMON runs it in fp16 (load_diffusion_cond does model.to(float16) after loading), so storing fp32 doubles the disk/transfer bytes for precision we discard. Add scripts/sa3/sa3_convert_fp16.py: rewrites a checkpoint casting only float32 tensors to fp16 (bf16 T5Gemma is already 2 bytes -> left alone; int/bool and the safetensors __metadata__ preserved). Supports an fp16 copy (default -fp16) or --in-place for the image-bake path. Numerically a no-op at runtime: today's path is fp16(fp32); pre-convert stores fp16(fp32) and load_state_dict upcasts fp16->fp32 (exact) before the same .to(fp16). Verified on the medium checkpoint: MAX_ABS_DIFF=0.0 across all 996 float tensors vs the fp32 load. DiT 9.22 GB -> 4.61 GB (checkpoint 9.8 -> 5.5 GiB, 44% smaller). tests/unit/test_sa3_fp16_convert.py locks the contract with synthetic safetensors (no GPU / weights in CI). Consuming this in the baked image (run --in-place before promote_warm) is Phase 1 (demon-public-demo) work, not part of this PR. Co-Authored-By: Claude Opus 4.8 (1M context) --- scripts/sa3/sa3_convert_fp16.py | 168 ++++++++++++++++++++++++++++ tests/unit/test_sa3_fp16_convert.py | 102 +++++++++++++++++ 2 files changed, 270 insertions(+) create mode 100644 scripts/sa3/sa3_convert_fp16.py create mode 100644 tests/unit/test_sa3_fp16_convert.py diff --git a/scripts/sa3/sa3_convert_fp16.py b/scripts/sa3/sa3_convert_fp16.py new file mode 100644 index 00000000..891dc1b7 --- /dev/null +++ b/scripts/sa3/sa3_convert_fp16.py @@ -0,0 +1,168 @@ +"""Convert an SA3 checkpoint's fp32 weights to fp16 on disk. + +The medium DiT ships as ~9 GB of **float32** safetensors, but DEMON runs +it in fp16 (``load_diffusion_cond(..., model_half=True)`` does +``model.to(torch.float16)`` after loading). Storing fp32 therefore costs +2x the disk/VRAM-transfer bytes for precision we immediately discard. + +This script rewrites the checkpoint with fp32 tensors cast to fp16, +which: + * halves the DiT file (~8.6 GiB -> ~4.3 GiB) — smaller Docker image + and less to read at load, and + * is numerically identical at runtime: today's path already does + fp16(fp32_weights); pre-converting yields fp16(fp32_weights) stored, + and load_state_dict upcasts fp16->fp32 (exact) before the same + ``.to(fp16)`` — so the resident weights are bit-for-bit the same. + +Only **float32** tensors are converted. bf16 tensors (the bundled +T5Gemma encoder) are left alone: bf16 is already 2 bytes, so converting +saves no space and would needlessly change its stored precision (runtime +still casts it to fp16 at load, exactly as before). int/bool tensors and +the safetensors ``__metadata__`` header are preserved verbatim. + +Usage: + # write an fp16 copy next to the source (default: -fp16/) + python scripts/sa3/sa3_convert_fp16.py --model-id medium + + # overwrite the canonical checkpoint in place (image-bake use case) + python scripts/sa3/sa3_convert_fp16.py --model-id medium --in-place +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import struct +import sys +from pathlib import Path + +_HERE = Path(__file__).resolve().parent +_REPO_ROOT = next(p for p in (_HERE, *_HERE.parents) if (p / "pyproject.toml").exists()) +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from acestep.engine.sa3_helpers import sa3_checkpoint_dir # noqa: E402 + + +def _header_dtypes(path: Path) -> set[str]: + """The set of tensor dtype strings in a safetensors file, read from + the JSON header only (no tensor data loaded).""" + with open(path, "rb") as fh: + (n,) = struct.unpack(" tuple[int, int]: + """Write ``dst`` from ``src`` with float32 tensors cast to fp16. + + Returns ``(converted, total)`` tensor counts. If the file has no + float32 tensors, it is copied byte-for-byte (no re-serialization) so + bf16/int payloads and header ordering are preserved exactly. + """ + if "F32" not in _header_dtypes(src): + shutil.copy2(src, dst) + return 0, 0 + + import torch # noqa: PLC0415 + from safetensors import safe_open # noqa: PLC0415 + from safetensors.torch import load_file, save_file # noqa: PLC0415 + + with safe_open(str(src), framework="pt") as f: + metadata = f.metadata() # preserved (may be None) + tensors = load_file(str(src)) + out: dict = {} + converted = 0 + for key, value in tensors.items(): + if value.dtype == torch.float32: + value = value.to(torch.float16) + converted += 1 + out[key] = value.contiguous() + save_file(out, str(dst), metadata=metadata) + return converted, len(tensors) + + +def convert_checkpoint(src_dir: Path, dst_dir: Path) -> None: + """Materialize an fp16 copy of an SA3 checkpoint dir at ``dst_dir``. + + Every ``*.safetensors`` has its float32 tensors cast to fp16; all + other files (model_config.json, tokenizer, t5gemma bf16 weights, …) + are copied verbatim. The HF ``.cache`` download dir is skipped. + """ + total_before = total_after = 0 + for src_path in sorted(src_dir.rglob("*")): + rel = src_path.relative_to(src_dir) + if ".cache" in rel.parts: + continue + dst_path = dst_dir / rel + if src_path.is_dir(): + dst_path.mkdir(parents=True, exist_ok=True) + continue + dst_path.parent.mkdir(parents=True, exist_ok=True) + if src_path.suffix == ".safetensors": + before = src_path.stat().st_size + converted, total = convert_or_copy_safetensors(src_path, dst_path) + after = dst_path.stat().st_size + total_before += before + total_after += after + print( + f" {rel}: {converted}/{total} fp32->fp16 " + f"{before / 1e9:.2f} GB -> {after / 1e9:.2f} GB" + ) + else: + shutil.copy2(src_path, dst_path) + if total_before: + print( + f"[convert] safetensors total: {total_before / 1e9:.2f} GB -> " + f"{total_after / 1e9:.2f} GB " + f"({100 * (1 - total_after / total_before):.0f}% smaller)" + ) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0]) + ap.add_argument("--model-id", default="medium", help='e.g. "medium", "small-music"') + ap.add_argument("--src", default=None, help="checkpoint dir (overrides --model-id)") + ap.add_argument("--out", default=None, help="output dir (default: -fp16)") + ap.add_argument( + "--in-place", action="store_true", + help="overwrite the source checkpoint's safetensors instead of copying", + ) + ap.add_argument("--force", action="store_true", help="overwrite an existing --out") + args = ap.parse_args() + + src = Path(args.src) if args.src else sa3_checkpoint_dir(args.model_id) + if not (src / "model.safetensors").is_file(): + print(f"[convert] no checkpoint at {src}", file=sys.stderr) + return 1 + + if args.in_place: + # Convert each safetensors to a temp sibling, then atomically + # replace — never leave a half-written weight file at the path + # the loader/preflight reads. + print(f"[convert] in-place: {src}") + for st in sorted(src.rglob("*.safetensors")): + if ".cache" in st.relative_to(src).parts: + continue + tmp = st.with_suffix(".safetensors.fp16.tmp") + converted, total = convert_or_copy_safetensors(st, tmp) + if converted: + os.replace(tmp, st) + print(f" {st.relative_to(src)}: {converted}/{total} fp32->fp16") + else: + tmp.unlink(missing_ok=True) + return 0 + + dst = Path(args.out) if args.out else src.parent / f"{src.name}-fp16" + if dst.exists() and not args.force: + print(f"[convert] {dst} exists (use --force to overwrite)", file=sys.stderr) + return 1 + print(f"[convert] {src} -> {dst}") + convert_checkpoint(src, dst) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit/test_sa3_fp16_convert.py b/tests/unit/test_sa3_fp16_convert.py new file mode 100644 index 00000000..0b056d14 --- /dev/null +++ b/tests/unit/test_sa3_fp16_convert.py @@ -0,0 +1,102 @@ +"""Tests for the SA3 fp16 checkpoint converter. + +Uses tiny synthetic safetensors rather than the real 9 GB checkpoint, so +it runs in CI without GPU or weights. The contract we lock: + * float32 tensors -> float16, values == the plain fp16 cast + * bf16 / int / bool tensors preserved verbatim (bf16 gains no space) + * safetensors __metadata__ and sidecar files preserved + * idempotent (a second pass converts nothing) +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_REPO_ROOT = Path(__file__).resolve().parents[2] +for _p in (str(_REPO_ROOT), str(_REPO_ROOT / "scripts" / "sa3")): + if _p not in sys.path: + sys.path.insert(0, _p) + +torch = pytest.importorskip("torch") +pytest.importorskip("safetensors") +from safetensors import safe_open # noqa: E402 +from safetensors.torch import load_file, save_file # noqa: E402 + +import sa3_convert_fp16 as conv # noqa: E402 + + +def _make_checkpoint(root: Path): + root.mkdir(parents=True, exist_ok=True) + dit = { + "w": torch.randn(8, 8, dtype=torch.float32), + "b": torch.randn(8, dtype=torch.float32), + "step": torch.tensor([3, 1, 4], dtype=torch.int64), # non-float + } + save_file(dit, str(root / "model.safetensors"), metadata={"format": "pt", "tag": "x"}) + (root / "model_config.json").write_text('{"model": {}}', encoding="utf-8") + enc = root / "t5gemma-b-b-ul2" + enc.mkdir(exist_ok=True) + # already-2-byte encoder weights: must be left untouched + save_file( + {"h": torch.randn(4, 4).to(torch.bfloat16)}, + str(enc / "model.safetensors"), + ) + (enc / "tokenizer.json").write_text("{}", encoding="utf-8") + return dit + + +def test_fp32_converted_bf16_and_int_preserved(tmp_path): + src = tmp_path / "src" + ref = _make_checkpoint(src) + dst = tmp_path / "dst" + + conv.convert_checkpoint(src, dst) + + out = load_file(str(dst / "model.safetensors")) + assert out["w"].dtype == torch.float16 + assert out["b"].dtype == torch.float16 + assert out["step"].dtype == torch.int64 # non-float kept + assert torch.equal(out["w"], ref["w"].to(torch.float16)) # faithful cast + assert torch.equal(out["step"], ref["step"]) + + enc = load_file(str(dst / "t5gemma-b-b-ul2" / "model.safetensors")) + assert enc["h"].dtype == torch.bfloat16 # bf16 untouched + + +def test_metadata_and_sidecars_preserved(tmp_path): + src = tmp_path / "src" + _make_checkpoint(src) + dst = tmp_path / "dst" + + conv.convert_checkpoint(src, dst) + + with safe_open(str(dst / "model.safetensors"), framework="pt") as f: + assert f.metadata().get("tag") == "x" + assert (dst / "model_config.json").read_text(encoding="utf-8") == '{"model": {}}' + assert (dst / "t5gemma-b-b-ul2" / "tokenizer.json").exists() + + +def test_bf16_only_file_is_copied_not_reserialized(tmp_path): + src = tmp_path / "enc.safetensors" + save_file({"h": torch.randn(4, 4).to(torch.bfloat16)}, str(src)) + dst = tmp_path / "enc_out.safetensors" + + converted, total = conv.convert_or_copy_safetensors(src, dst) + + assert (converted, total) == (0, 0) # copy path + assert src.read_bytes() == dst.read_bytes() # byte-identical + + +def test_idempotent(tmp_path): + src = tmp_path / "src" + _make_checkpoint(src) + dst1 = tmp_path / "dst1" + conv.convert_checkpoint(src, dst1) + # second pass over an already-fp16 dir converts nothing + converted, total = conv.convert_or_copy_safetensors( + dst1 / "model.safetensors", tmp_path / "again.safetensors" + ) + assert converted == 0