scripts/70_convert_checkpoint.py

#!/usr/bin/env python
"""
70_convert_checkpoint.py — Convert a DeepSpeed pipeline-parallel checkpoint
to an Anima-loadable safetensors file.

Why this exists: diffusion-pipe's mid-training checkpoints are deepspeed
pipeline-parallel format (one .pt file per pipeline "layer"). The clean
.safetensors format only appears at epoch boundaries via the [save] phase,
and our epoch boundary (step ~7550 for full 4-variant pass) is past our
max_steps cap. To run a mid-training checkpoint in ComfyUI / inference for
comparison purposes, we need to consolidate the layer .pt files into one
.safetensors with the same key structure as the original Anima Preview3
base weights.

The mapping was reverse-engineered by inspecting the structure:

  DeepSpeed file              Contents               Anima safetensors prefix
  ---------------             --------               ------------------------
  layer_00-model_states.pt    x_embedder.*,          net.*  (no rename)
                              t_embedder.*,
                              t_embedding_norm.*
  layer_01-model_states.pt    (empty)                skip
  layer_02-model_states.pt    block.self_attn.*,     net.blocks.0.*  (drop "block.")
                              block.cross_attn.*,
                              block.mlp.*,
                              block.adaln_modulation_*
  layer_03-model_states.pt    same shape             net.blocks.1.*
  ...
  layer_29-model_states.pt    same shape             net.blocks.27.*
  layer_30-model_states.pt    final_layer.*          net.* (no rename)
  layer_31-model_states.pt    (empty, if present)    skip

So the general rules:
- layer_00 keys: prepend "net."
- layer_NN where N in [2..N_blocks+1]: rename "block." → "net.blocks.<N-2>."
- last meaningful layer (containing final_layer.*): prepend "net."
- empty layers: skip

After conversion, the script verifies:
  1. Output key set matches Anima base key set exactly (set equality)
  2. Per-key tensor shapes match between converted output and Anima base

Usage:
  python 70_convert_checkpoint.py \\
      --deepspeed-dir /opt/local/outputs/anima_finish_v1/<run>/global_step1670 \\
      --reference-safetensors /opt/local/models/anima/split_files/diffusion_models/anima-preview3-base.safetensors \\
      --output /workspace/checkpoints/anima_finish_v1_step1670.safetensors

Exit codes:
  0 = converted + verified
  1 = converted but verification failed
  2 = no valid checkpoint files found
  3 = ref file missing
"""
from __future__ import annotations

import argparse
import sys
import time
from collections import OrderedDict
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import save_file


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    p.add_argument("--deepspeed-dir", required=True, type=Path,
                   help="Directory containing layer_NN-model_states.pt files.")
    p.add_argument("--reference-safetensors", required=True, type=Path,
                   help="Original Anima safetensors for key/shape verification.")
    p.add_argument("--output", required=True, type=Path,
                   help="Path to write the converted safetensors.")
    p.add_argument("--dtype", default="bfloat16",
                   choices=["bfloat16", "float16", "float32"],
                   help="Output tensor dtype. Default bfloat16 matches Anima.")
    p.add_argument("--skip-verify", action="store_true",
                   help="Skip key/shape verification (faster, but unsafe).")
    p.add_argument("--merge-frozen-from-ref", action="store_true",
                   help="For any keys in the reference but missing from "
                        "the deepspeed checkpoint, copy them from the "
                        "reference safetensors. Use when training has "
                        "frozen modules (e.g. llm_adapter_lr=0) — those "
                        "weights aren't saved in checkpoints but are "
                        "needed for inference. With this flag, the output "
                        "is a complete model = trained_modules ∪ frozen_from_ref.")
    return p.parse_args()


def load_reference_keys(ref_path: Path) -> dict[str, torch.Size]:
    """Return {key: shape} from the reference safetensors."""
    ref: dict[str, torch.Size] = {}
    with safe_open(str(ref_path), framework="pt", device="cpu") as f:
        for k in f.keys():
            t = f.get_tensor(k)
            ref[k] = t.shape
    return ref


def discover_layers(ckpt_dir: Path) -> list[Path]:
    """Return sorted list of layer_NN-model_states.pt paths."""
    files = sorted(ckpt_dir.glob("layer_*-model_states.pt"))
    if not files:
        print(f"ERROR: no layer_*-model_states.pt under {ckpt_dir}", file=sys.stderr)
        sys.exit(2)
    return files


def load_layer_state(path: Path) -> dict[str, torch.Tensor]:
    """Load one layer file. Returns its state_dict (empty dict if empty)."""
    sd = torch.load(str(path), map_location="cpu", weights_only=False)
    if not isinstance(sd, dict):
        print(f"WARN  unexpected type in {path.name}: {type(sd).__name__}",
              file=sys.stderr)
        return {}
    # Filter to only tensors (skip metadata if any).
    return {k: v for k, v in sd.items() if isinstance(v, torch.Tensor)}


def convert(ckpt_dir: Path) -> OrderedDict[str, torch.Tensor]:
    """
    Walk all layer .pt files in order, remap keys to Anima's namespace.

    Strategy:
      - Bucket layers into [embedding, block_0, block_1, ..., final].
      - Empty layers are skipped (passthrough).
      - The FIRST non-empty layer with x_embedder/t_embedder/etc keys becomes
        the embedding layer (prepend "net.").
      - The LAST non-empty layer with final_layer.* keys becomes the final
        layer (prepend "net.").
      - Everything in between with "block.*" keys is a transformer block,
        rename to "net.blocks.<i>." where i counts up from 0 across them.
    """
    layer_files = discover_layers(ckpt_dir)
    print(f"found {len(layer_files)} layer files in {ckpt_dir.name}")

    # First pass: load every layer, classify it.
    loaded: list[tuple[Path, dict[str, torch.Tensor]]] = []
    for f in layer_files:
        sd = load_layer_state(f)
        loaded.append((f, sd))

    # Classify
    embedding_layers: list[dict[str, torch.Tensor]] = []
    block_layers: list[dict[str, torch.Tensor]] = []
    final_layers: list[dict[str, torch.Tensor]] = []
    empty_count = 0

    for f, sd in loaded:
        if not sd:
            empty_count += 1
            continue
        # Determine type by inspecting top-level keys
        any_block = any(k.startswith("block.") for k in sd.keys())
        any_final = any(k.startswith("final_layer.") for k in sd.keys())
        any_embed = any(k.startswith(p) for k in sd.keys()
                        for p in ("x_embedder.", "t_embedder.", "t_embedding_norm."))
        if any_block:
            block_layers.append(sd)
        elif any_final:
            final_layers.append(sd)
        elif any_embed:
            embedding_layers.append(sd)
        else:
            print(f"WARN unclassifiable layer {f.name}: top keys = "
                  f"{list(sd.keys())[:3]}", file=sys.stderr)

    print(f"  embedding layers: {len(embedding_layers)}")
    print(f"  block layers: {len(block_layers)}")
    print(f"  final layers: {len(final_layers)}")
    print(f"  empty (passthrough) layers: {empty_count}")

    if len(embedding_layers) != 1:
        print(f"WARN expected 1 embedding layer, found {len(embedding_layers)}",
              file=sys.stderr)
    if len(final_layers) != 1:
        print(f"WARN expected 1 final layer, found {len(final_layers)}",
              file=sys.stderr)

    # Build the output state dict with Anima's "net." prefix.
    out: OrderedDict[str, torch.Tensor] = OrderedDict()

    # Embedding (prepend "net.")
    for sd in embedding_layers:
        for k, v in sd.items():
            out[f"net.{k}"] = v

    # Blocks (rename "block." → "net.blocks.<i>.")
    for i, sd in enumerate(block_layers):
        prefix = f"net.blocks.{i}."
        for k, v in sd.items():
            if k.startswith("block."):
                out[prefix + k[len("block."):]] = v
            else:
                print(f"WARN block {i} has unexpected key '{k}'", file=sys.stderr)
                out[prefix + k] = v

    # Final (prepend "net.")
    for sd in final_layers:
        for k, v in sd.items():
            out[f"net.{k}"] = v

    return out


def verify(out: dict[str, torch.Tensor], ref: dict[str, torch.Size]) -> bool:
    """Return True iff converted output matches reference (keys + shapes)."""
    ok = True

    out_keys = set(out.keys())
    ref_keys = set(ref.keys())

    missing_in_out = ref_keys - out_keys
    extra_in_out = out_keys - ref_keys

    if missing_in_out:
        print(f"FAIL  {len(missing_in_out)} keys missing in converted output:",
              file=sys.stderr)
        for k in sorted(missing_in_out)[:10]:
            print(f"  - {k}", file=sys.stderr)
        if len(missing_in_out) > 10:
            print(f"  ... and {len(missing_in_out) - 10} more", file=sys.stderr)
        ok = False

    if extra_in_out:
        print(f"FAIL  {len(extra_in_out)} keys in converted output that aren't in reference:",
              file=sys.stderr)
        for k in sorted(extra_in_out)[:10]:
            print(f"  + {k}", file=sys.stderr)
        if len(extra_in_out) > 10:
            print(f"  ... and {len(extra_in_out) - 10} more", file=sys.stderr)
        ok = False

    # Shape verification on overlapping keys
    shape_mismatches: list[tuple[str, tuple, tuple]] = []
    for k in out_keys & ref_keys:
        out_shape = tuple(out[k].shape)
        ref_shape = tuple(ref[k])
        if out_shape != ref_shape:
            shape_mismatches.append((k, out_shape, ref_shape))

    if shape_mismatches:
        print(f"FAIL  {len(shape_mismatches)} shape mismatches:",
              file=sys.stderr)
        for k, o, r in shape_mismatches[:10]:
            print(f"  {k}: converted={o} vs reference={r}", file=sys.stderr)
        if len(shape_mismatches) > 10:
            print(f"  ... and {len(shape_mismatches) - 10} more", file=sys.stderr)
        ok = False

    if ok:
        print(f"OK   {len(out_keys)} keys match reference exactly")
        print(f"OK   all tensor shapes match")
    return ok


def main() -> int:
    args = parse_args()
    if not args.deepspeed_dir.is_dir():
        print(f"--deepspeed-dir not a directory: {args.deepspeed_dir}",
              file=sys.stderr)
        return 2
    if not args.reference_safetensors.is_file():
        print(f"--reference-safetensors missing: {args.reference_safetensors}",
              file=sys.stderr)
        return 3

    target_dtype = {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }[args.dtype]

    print(f"loading reference key/shape index from {args.reference_safetensors.name} ...")
    t0 = time.monotonic()
    ref = load_reference_keys(args.reference_safetensors)
    print(f"  {len(ref)} keys in {time.monotonic() - t0:.1f}s")

    print()
    print(f"converting {args.deepspeed_dir.name} ...")
    t0 = time.monotonic()
    out = convert(args.deepspeed_dir)
    print(f"  converted {len(out)} keys in {time.monotonic() - t0:.1f}s")

    # Merge in frozen modules from the reference (e.g. llm_adapter when
    # llm_adapter_lr=0 was set in training). DeepSpeed only saves trained
    # parameters, so frozen ones are missing from the checkpoint — but
    # they're identical to base Anima by construction (no gradient updates).
    if args.merge_frozen_from_ref:
        missing_in_out = set(ref.keys()) - set(out.keys())
        if missing_in_out:
            print()
            print(f"merging {len(missing_in_out)} frozen keys from reference "
                  f"(--merge-frozen-from-ref) ...")
            top_prefixes = set()
            for k in missing_in_out:
                top_prefixes.add(".".join(k.split(".")[:2]))
            print(f"  top-level prefixes being merged: {sorted(top_prefixes)}")
            with safe_open(str(args.reference_safetensors),
                           framework="pt", device="cpu") as f:
                for k in missing_in_out:
                    out[k] = f.get_tensor(k)
            print(f"  merged. output now has {len(out)} keys.")

    # Cast all tensors to target dtype
    if target_dtype != torch.float32:
        print(f"casting to {args.dtype} ...")
        out = OrderedDict((k, v.to(target_dtype)) for k, v in out.items())

    # Verify
    if not args.skip_verify:
        print()
        print("verifying ...")
        if not verify(out, ref):
            print("verification FAILED — not writing output", file=sys.stderr)
            return 1

    # Save
    args.output.parent.mkdir(parents=True, exist_ok=True)
    print()
    print(f"writing {args.output} ...")
    t0 = time.monotonic()
    # Make all tensors contiguous (safetensors requires it)
    out = OrderedDict((k, v.contiguous()) for k, v in out.items())
    save_file(out, str(args.output))
    sz = args.output.stat().st_size / (1024 ** 3)
    print(f"  wrote {sz:.2f} GB in {time.monotonic() - t0:.1f}s")
    return 0


if __name__ == "__main__":
    sys.exit(main())