runtime/packed_artifact.py

#!/usr/bin/env python3

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import sys
from collections import OrderedDict
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable, Iterable

import safetensors
import torch
import torch.nn as nn
from safetensors import safe_open
from safetensors.torch import save_file


FORMAT_VERSION = 1
ARTIFACT_KIND = "mage_flow_transformer_mlp_nvfp4_resident_v1"
CANONICAL_SAFETENSORS_METADATA = {
    "mage_nvfp4_contract": (
        f"{ARTIFACT_KIND};format_version={FORMAT_VERSION}"
    )
}
LEGACY_SAFETENSORS_METADATA = {
    "artifact_kind": ARTIFACT_KIND,
    "format_version": str(FORMAT_VERSION),
}
FP4_BLOCK_ELEMENTS = 16
SCALE_TILE_OUTER = 128
SCALE_TILE_INNER = 4
FP4_E2M1_MAX = 6.0
FP4_TENSOR_SCALE_MAX = 448.0
NVFP4_TENSOR_SCALE_DENOMINATOR = FP4_E2M1_MAX * FP4_TENSOR_SCALE_MAX
TARGET_DEPTH = 12

RELEASE_ROOT = Path(__file__).resolve().parents[1]
PROJECT_ROOT = RELEASE_ROOT
MAGE_ROOT = RELEASE_ROOT / "vendor"
RESIDENT_PYTHON_ROOT = RELEASE_ROOT / "runtime"
RESIDENT_SOURCE = RESIDENT_PYTHON_ROOT / "nvfp4_linear.cu"
RESIDENT_LIBRARY = RESIDENT_PYTHON_ROOT / "libmage_nvfp4_linear.so"
ARTIFACT_SCRIPT_PATH = Path(__file__).resolve()
_NATIVE_PACKER = None


class PackedArtifactError(RuntimeError):
    pass


@dataclass(frozen=True)
class ScaleLayout:
    inner_dim: int
    outer_tiles: int
    bytes: int


@dataclass(frozen=True)
class TargetSpec:
    module_key: str
    weight_key: str
    bias_key: str
    artifact_weight_key: str
    artifact_scale_key: str
    artifact_tensor_scale_key: str
    artifact_bias_key: str


def fail(message: str) -> None:
    raise PackedArtifactError(message)


def round_up(value: int, multiple: int) -> int:
    return ((value + multiple - 1) // multiple) * multiple


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while True:
            chunk = handle.read(1 << 20)
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest()


def sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def tensor_bytes(tensor: torch.Tensor) -> bytes:
    if not tensor.is_contiguous():
        tensor = tensor.contiguous()
    return tensor.view(torch.uint8).cpu().numpy().tobytes()


def sha256_tensor(tensor: torch.Tensor) -> str:
    return sha256_bytes(tensor_bytes(tensor))


def fsync_directory(path: Path) -> None:
    descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)


def fsync_file(path: Path) -> None:
    descriptor = os.open(path, os.O_RDONLY)
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)


def write_bytes_once(path: Path, payload: bytes) -> None:
    with path.open("xb") as handle:
        handle.write(payload)
        handle.flush()
        os.fsync(handle.fileno())
    fsync_directory(path.parent)


def host_scale_offset(outer: int, inner_scale: int, scale_inner_dim: int) -> int:
    outer_tile = outer // SCALE_TILE_OUTER
    local_outer = outer % SCALE_TILE_OUTER
    local_inner = inner_scale % SCALE_TILE_INNER
    inner_tile_start = inner_scale - local_inner
    tile_base = (inner_tile_start + outer_tile * scale_inner_dim) * SCALE_TILE_OUTER
    return tile_base + (local_outer % 32) * 16 + (local_outer // 32) * 4 + local_inner


def make_scale_layout(rows_k: int, outer_columns: int) -> ScaleLayout:
    if rows_k <= 0 or outer_columns <= 0:
        fail("scale layout requires positive rows_k and outer_columns")
    if rows_k % FP4_BLOCK_ELEMENTS:
        fail(
            f"scale layout requires K divisible by {FP4_BLOCK_ELEMENTS}; "
            f"got K={rows_k}"
        )
    inner_dim = round_up(rows_k // FP4_BLOCK_ELEMENTS, SCALE_TILE_INNER)
    outer_tiles = (outer_columns + SCALE_TILE_OUTER - 1) // SCALE_TILE_OUTER
    return ScaleLayout(
        inner_dim=inner_dim,
        outer_tiles=outer_tiles,
        bytes=outer_tiles * inner_dim * SCALE_TILE_OUTER,
    )


def build_target_specs(depth: int) -> list[TargetSpec]:
    if depth != TARGET_DEPTH:
        fail(
            f"this artifact format is pinned to exactly {TARGET_DEPTH} transformer blocks; "
            f"config reported depth={depth}"
        )
    specs: list[TargetSpec] = []
    suffixes = (
        "img_mlp.net.0.proj",
        "img_mlp.net.2",
        "txt_mlp.net.0.proj",
        "txt_mlp.net.2",
    )
    for index in range(depth):
        for suffix in suffixes:
            module_key = f"transformer_blocks.{index}.{suffix}"
            specs.append(
                TargetSpec(
                    module_key=module_key,
                    weight_key=f"{module_key}.weight",
                    bias_key=f"{module_key}.bias",
                    artifact_weight_key=f"targets.{module_key}.packed_weight_e2m1",
                    artifact_scale_key=f"targets.{module_key}.packed_scales_ue4m3",
                    artifact_tensor_scale_key=f"targets.{module_key}.weight_tensor_scale",
                    artifact_bias_key=f"targets.{module_key}.bias_bf16",
                )
            )
    return specs


def decode_fp4_e2m1(raw: int) -> float:
    sign = -1.0 if (raw & 0x8) else 1.0
    magnitude = raw & 0x7
    table = (
        0.0,
        0.5,
        1.0,
        1.5,
        2.0,
        3.0,
        4.0,
        6.0,
    )
    return sign * table[magnitude]


def encode_fp4_e2m1(value: float) -> int:
    candidates = [decode_fp4_e2m1(code) for code in range(16)]
    best_code = 0
    best_error = math.inf
    for code, candidate in enumerate(candidates):
        error = abs(candidate - value)
        if error < best_error or (error == best_error and (code & 1) == 0 and (best_code & 1) == 1):
            best_error = error
            best_code = code
    return best_code


def decode_fp8_e4m3(raw: int) -> float:
    sign = -1.0 if (raw & 0x80) else 1.0
    exponent = (raw >> 3) & 0x0F
    mantissa = raw & 0x07
    if exponent == 0:
        if mantissa == 0:
            return 0.0 * sign
        return sign * (mantissa / 8.0) * (2.0 ** -6)
    if exponent == 0x0F and mantissa == 0x07:
        return math.nan
    return sign * (1.0 + mantissa / 8.0) * (2.0 ** (exponent - 7))


def _build_positive_e4m3_table() -> list[tuple[int, float]]:
    table: list[tuple[int, float]] = []
    for raw in range(0x80):
        value = decode_fp8_e4m3(raw)
        if math.isnan(value) or value < 0.0:
            continue
        table.append((raw, value))
    table.sort(key=lambda item: (item[1], item[0]))
    return table


POSITIVE_E4M3_TABLE = _build_positive_e4m3_table()


def encode_fp8_e4m3_satfinite(value: float) -> int:
    if value <= 0.0:
        return 0
    finite_values = [item for item in POSITIVE_E4M3_TABLE if item[1] <= FP4_TENSOR_SCALE_MAX]
    best_raw = finite_values[-1][0]
    best_error = math.inf
    for raw, candidate in finite_values:
        error = abs(candidate - value)
        if error < best_error or (error == best_error and (raw & 1) == 0 and (best_raw & 1) == 1):
            best_error = error
            best_raw = raw
    return best_raw


def host_tensor_scale_from_amax(amax: float) -> float:
    return 1.0 if amax == 0.0 else amax / NVFP4_TENSOR_SCALE_DENOMINATOR


def native_packer():
    global _NATIVE_PACKER
    if _NATIVE_PACKER is None:
        if str(RESIDENT_PYTHON_ROOT) not in sys.path:
            sys.path.insert(0, str(RESIDENT_PYTHON_ROOT))
        from packed_nvfp4_linear import NativeNvfp4Library

        _NATIVE_PACKER = NativeNvfp4Library(RESIDENT_LIBRARY)
    return _NATIVE_PACKER


def pack_weight_tensor(weight_nk_bf16: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]:
    if weight_nk_bf16.dtype != torch.bfloat16 or weight_nk_bf16.device.type != "cpu":
        fail("weight packer expects a CPU bfloat16 tensor")
    if weight_nk_bf16.ndim != 2:
        fail("weight packer expects a 2D [N,K] weight tensor")
    if not weight_nk_bf16.is_contiguous():
        weight_nk_bf16 = weight_nk_bf16.contiguous()

    columns_n, rows_k = weight_nk_bf16.shape
    if rows_k % 32 != 0 or columns_n % 8 != 0:
        fail(
            f"native NVFP4 packing requires K%32==0 and N%8==0; got N={columns_n} K={rows_k}"
        )

    amax = float(weight_nk_bf16.float().abs().max().item())
    packed_fp4, packed_scales, tensor_scale = native_packer().pack_weight(
        weight_nk_bf16
    )
    return packed_fp4, packed_scales, tensor_scale.reshape(1), amax


def load_transformer_config(source_repo: Path) -> dict:
    config_path = source_repo / "transformer" / "config.json"
    if not config_path.exists():
        fail(f"missing transformer config: {config_path}")
    return json.loads(config_path.read_text())


def source_transformer_checkpoint(source_repo: Path) -> Path:
    checkpoint = source_repo / "transformer" / "diffusion_pytorch_model.safetensors"
    if not checkpoint.exists():
        fail(f"missing transformer checkpoint: {checkpoint}")
    return checkpoint


def import_mage_transformer_symbols() -> tuple[type[nn.Module], object]:
    if str(MAGE_ROOT) not in sys.path:
        sys.path.insert(0, str(MAGE_ROOT))
    try:
        from mage_flow.models.mage_flow import MageFlow, MageFlowParams
    except Exception as exc:
        fail(f"failed to import local Mage transformer sources from {MAGE_ROOT}: {exc}")
    return MageFlow, MageFlowParams


class PackedNvfp4LinearArtifactModule(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        packed_weight_e2m1: torch.Tensor,
        packed_scales_ue4m3: torch.Tensor,
        weight_tensor_scale: torch.Tensor,
        bias_bf16: torch.Tensor | None,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.register_buffer("packed_weight_e2m1", packed_weight_e2m1.contiguous())
        self.register_buffer("packed_scales_ue4m3", packed_scales_ue4m3.contiguous())
        self.register_buffer("weight_tensor_scale", weight_tensor_scale.contiguous())
        if bias_bf16 is None:
            self.bias_bf16 = None
        else:
            self.register_buffer("bias_bf16", bias_bf16.contiguous())

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        raise RuntimeError(
            "PackedNvfp4LinearArtifactModule is an artifact-only placeholder. "
            "Attach the resident CUDA runtime before calling forward()."
        )

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"packed_weight_bytes={self.packed_weight_e2m1.numel()}, "
            f"packed_scale_bytes={self.packed_scales_ue4m3.numel()}, "
            f"has_bias={self.bias_bf16 is not None}"
        )


def instantiate_mage_transformer_on_meta(source_repo: Path) -> nn.Module:
    config = load_transformer_config(source_repo)
    MageFlow, MageFlowParams = import_mage_transformer_symbols()
    structure = {
        key: value
        for key, value in config.items()
        if key
        not in {
            "_class_name",
            "txt_max_length",
            "max_sequence_length",
            "param_dtype",
            "packing",
            "schedule_mode",
            "static_shift",
            "use_time_shift",
            "rope_type",
            "apply_text_rotary_emb",
            "mlp_ratio",
            "depth_single_blocks",
            "theta",
            "qkv_bias",
            "guidance_embed",
            "vec_in_dim",
            "vec_type",
            "time_type",
            "double_block_type",
            "quantization_config",
        }
    }
    with torch.device("meta"):
        model = MageFlow(MageFlowParams(**structure))
    return model


def unregistered_meta_tensor_attribute_names(model: nn.Module) -> list[str]:
    """Find direct tensor attributes that PyTorch's parameter/buffer walk misses."""
    names: list[str] = []
    for module_name, module in model.named_modules():
        registered_names = set(module._parameters) | set(module._buffers)
        for attribute_name, value in vars(module).items():
            if attribute_name in registered_names:
                continue
            if isinstance(value, torch.Tensor) and value.is_meta:
                prefix = f"{module_name}." if module_name else ""
                names.append(f"{prefix}{attribute_name}")
    return sorted(names)


def materialize_mage_rope_tensor_attributes(model: nn.Module) -> list[str]:
    """Rebuild Mage's intentionally unregistered complex RoPE tensors on CPU."""
    before = unregistered_meta_tensor_attribute_names(model)
    expected = ["pos_embed.neg_freqs", "pos_embed.pos_freqs"]
    if before != expected:
        fail(
            "unexpected unregistered meta tensor attributes before RoPE "
            f"materialization: {before}"
        )

    rope = model.get_submodule("pos_embed")
    rope_type = type(rope)
    with torch.device("cpu"):
        materialized = rope_type(
            theta=rope.theta,
            axes_dim=list(rope.axes_dim),
            scale_rope=rope.scale_rope,
        )
    rope.pos_freqs = materialized.pos_freqs
    rope.neg_freqs = materialized.neg_freqs
    rope.video_freq_cache = {}

    remaining = unregistered_meta_tensor_attribute_names(model)
    if remaining:
        fail(
            "unresolved unregistered meta tensor attributes after RoPE "
            f"materialization: {remaining}"
        )
    return before


def set_child_module(root: nn.Module, dotted_path: str, module: nn.Module) -> None:
    parent_path, _, child_name = dotted_path.rpartition(".")
    parent = root.get_submodule(parent_path) if parent_path else root
    if child_name.isdigit() and isinstance(parent, (nn.Sequential, nn.ModuleList)):
        parent[int(child_name)] = module
    else:
        setattr(parent, child_name, module)


def replace_targets_with_artifact_modules(
    model: nn.Module,
    artifact_path: Path,
    target_specs: Iterable[TargetSpec],
) -> None:
    with safe_open(artifact_path, framework="pt", device="cpu") as handle:
        for spec in target_specs:
            packed_weight = handle.get_tensor(spec.artifact_weight_key)
            packed_scales = handle.get_tensor(spec.artifact_scale_key)
            weight_tensor_scale = handle.get_tensor(spec.artifact_tensor_scale_key)
            bias = handle.get_tensor(spec.artifact_bias_key)
            original = model.get_submodule(spec.module_key)
            if not isinstance(original, nn.Linear):
                fail(f"expected target module {spec.module_key} to be nn.Linear")
            replacement = PackedNvfp4LinearArtifactModule(
                in_features=int(original.in_features),
                out_features=int(original.out_features),
                packed_weight_e2m1=packed_weight,
                packed_scales_ue4m3=packed_scales,
                weight_tensor_scale=weight_tensor_scale,
                bias_bf16=bias,
            )
            set_child_module(model, spec.module_key, replacement)


def replace_targets_with_resident_modules(
    model: nn.Module,
    artifact_path: Path,
    target_specs: Iterable[TargetSpec],
    device: torch.device,
) -> None:
    if device.type != "cuda":
        fail("resident runtime modules require a CUDA destination")
    if str(RESIDENT_PYTHON_ROOT) not in sys.path:
        sys.path.insert(0, str(RESIDENT_PYTHON_ROOT))
    from torch_ops import PackedNvfp4LinearOp

    _replace_targets_with_registered_modules(
        model,
        artifact_path,
        target_specs,
        device,
        PackedNvfp4LinearOp,
    )


def replace_targets_with_native_resident_modules(
    model: nn.Module,
    artifact_path: Path,
    target_specs: Iterable[TargetSpec],
    device: torch.device,
) -> None:
    if device.type != "cuda":
        fail("native resident runtime modules require a CUDA destination")
    if str(RESIDENT_PYTHON_ROOT) not in sys.path:
        sys.path.insert(0, str(RESIDENT_PYTHON_ROOT))
    from torch_ops_native import (
        PackedNvfp4LinearNativeOp,
        initialize_native_sm120_op,
    )

    if not initialize_native_sm120_op(allow_python_schema_fallback=False):
        fail("compiled native resident torch op did not load")
    _replace_targets_with_registered_modules(
        model,
        artifact_path,
        target_specs,
        device,
        PackedNvfp4LinearNativeOp,
    )


def _replace_targets_with_registered_modules(
    model: nn.Module,
    artifact_path: Path,
    target_specs: Iterable[TargetSpec],
    device: torch.device,
    module_cls: type[nn.Module],
) -> None:
    with safe_open(artifact_path, framework="pt", device="cpu") as handle:
        for spec in target_specs:
            original = model.get_submodule(spec.module_key)
            if not isinstance(original, nn.Linear):
                fail(f"expected target module {spec.module_key} to be nn.Linear")
            replacement = module_cls(
                in_features=int(original.in_features),
                out_features=int(original.out_features),
                packed_weight=handle.get_tensor(spec.artifact_weight_key).to(device),
                weight_scales=handle.get_tensor(spec.artifact_scale_key).to(device),
                weight_scale=handle.get_tensor(
                    spec.artifact_tensor_scale_key
                ).to(device),
                bias=handle.get_tensor(spec.artifact_bias_key).to(device),
            )
            set_child_module(model, spec.module_key, replacement)


def assign_tensor_by_name(model: nn.Module, key: str, tensor: torch.Tensor) -> None:
    if "." not in key:
        parent = model
        leaf = key
    else:
        parent_path, _, leaf = key.rpartition(".")
        parent = model.get_submodule(parent_path)
    if leaf in parent._parameters:
        requires_grad = parent._parameters[leaf].requires_grad
        parent._parameters[leaf] = nn.Parameter(tensor, requires_grad=requires_grad)
        return
    if leaf in parent._buffers:
        parent._buffers[leaf] = tensor
        return
    fail(f"destination key {key} was neither a parameter nor a buffer")


def pack_artifact(source_repo: Path, output_dir: Path) -> Path:
    source_repo = source_repo.resolve()
    output_dir = output_dir.resolve()
    if output_dir.exists():
        fail(f"output directory already exists: {output_dir}")
    output_dir.mkdir(parents=True, exist_ok=False)

    config = load_transformer_config(source_repo)
    checkpoint_path = source_transformer_checkpoint(source_repo)
    target_specs = build_target_specs(int(config["depth"]))
    target_weight_keys = {spec.weight_key for spec in target_specs}
    target_bias_keys = {spec.bias_key for spec in target_specs}
    target_keys = target_weight_keys | target_bias_keys

    artifact_tensors: OrderedDict[str, torch.Tensor] = OrderedDict()
    target_metadata: list[dict] = []

    with safe_open(checkpoint_path, framework="pt", device="cpu") as handle:
        source_keys = list(handle.keys())
        source_key_set = set(source_keys)
        missing = sorted(target_keys - source_key_set)
        if missing:
            fail(f"source checkpoint is missing {len(missing)} target tensors, first={missing[0]}")

        for spec in target_specs:
            weight = handle.get_tensor(spec.weight_key)
            bias = handle.get_tensor(spec.bias_key)
            if weight.dtype != torch.bfloat16:
                fail(f"{spec.weight_key} expected bfloat16, found {weight.dtype}")
            if bias.dtype != torch.bfloat16:
                fail(f"{spec.bias_key} expected bfloat16, found {bias.dtype}")
            packed_weight, packed_scales, weight_tensor_scale, global_amax = pack_weight_tensor(weight)
            scale_layout = make_scale_layout(weight.shape[1], weight.shape[0])
            artifact_tensors[spec.artifact_bias_key] = bias.contiguous()
            artifact_tensors[spec.artifact_scale_key] = packed_scales
            artifact_tensors[spec.artifact_tensor_scale_key] = weight_tensor_scale
            artifact_tensors[spec.artifact_weight_key] = packed_weight
            target_metadata.append(
                {
                    "module_key": spec.module_key,
                    "weight_key": spec.weight_key,
                    "bias_key": spec.bias_key,
                    "weight_shape": list(weight.shape),
                    "bias_shape": list(bias.shape),
                    "weight_tensor_scale_key": spec.artifact_tensor_scale_key,
                    "artifact_weight_key": spec.artifact_weight_key,
                    "artifact_scale_key": spec.artifact_scale_key,
                    "artifact_bias_key": spec.artifact_bias_key,
                    "weight_tensor_scale": float(weight_tensor_scale.item()),
                    "weight_global_amax": global_amax,
                    "packed_weight_bytes": int(packed_weight.numel()),
                    "packed_scale_bytes": int(packed_scales.numel()),
                    "scale_layout": asdict(scale_layout),
                    "source_weight_sha256": sha256_tensor(weight),
                    "source_bias_sha256": sha256_tensor(bias),
                }
            )

    non_target_keys = sorted(set(source_keys) - target_keys)
    artifact_path = output_dir / "packed_transformer.safetensors"
    save_file(
        OrderedDict(sorted(artifact_tensors.items())),
        artifact_path,
        metadata=CANONICAL_SAFETENSORS_METADATA,
    )
    fsync_file(artifact_path)
    fsync_directory(output_dir)

    library_hashes = {
        "artifact_script_sha256": sha256_file(ARTIFACT_SCRIPT_PATH),
        "resident_source_sha256": sha256_file(RESIDENT_SOURCE),
        "resident_library_sha256": sha256_file(RESIDENT_LIBRARY),
        "mage_flow_py_sha256": sha256_file(MAGE_ROOT / "mage_flow" / "models" / "mage_flow.py"),
        "mage_layers_py_sha256": sha256_file(MAGE_ROOT / "mage_flow" / "models" / "modules" / "mage_layers.py"),
        "pipeline_py_sha256": sha256_file(MAGE_ROOT / "mage_flow" / "pipeline.py"),
    }

    metadata = OrderedDict(
        (
            ("format_version", FORMAT_VERSION),
            ("artifact_kind", ARTIFACT_KIND),
            ("created_utc", datetime.now(timezone.utc).isoformat()),
            (
                "container",
                {
                    "format": "safetensors",
                    "header_metadata": CANONICAL_SAFETENSORS_METADATA,
                    "header_encoding": (
                        "single deterministic contract key; legacy two-key "
                        "draft headers remain readable"
                    ),
                },
            ),
            (
                "source",
                OrderedDict(
                    (
                        ("transformer_config_path", str(source_repo / "transformer" / "config.json")),
                        ("transformer_checkpoint_path", str(checkpoint_path)),
                        ("transformer_config_sha256", sha256_file(source_repo / "transformer" / "config.json")),
                        ("transformer_checkpoint_sha256", sha256_file(checkpoint_path)),
                    )
                ),
            ),
            ("library_hashes", library_hashes),
            (
                "environment",
                {
                    "python_version": sys.version,
                    "torch_version": torch.__version__,
                    "safetensors_version": safetensors.__version__,
                },
            ),
            (
                "model",
                {
                    "depth": int(config["depth"]),
                    "hidden_size": int(config["hidden_size"]),
                    "num_heads": int(config["num_heads"]),
                    "context_in_dim": int(config["context_in_dim"]),
                    "in_channels": int(config["in_channels"]),
                    "out_channels": int(config["out_channels"]),
                    "patch_size": int(config["patch_size"]),
                },
            ),
            (
                "quantization",
                {
                    "format": "nvfp4_two_level",
                    "block_elements": FP4_BLOCK_ELEMENTS,
                    "scale_tile_outer": SCALE_TILE_OUTER,
                    "scale_tile_inner": SCALE_TILE_INNER,
                    "bias_policy": "artifact_bfloat16",
                },
            ),
            ("targets", target_metadata),
            ("non_target_keys", non_target_keys),
        )
    )
    write_bytes_once(
        output_dir / "metadata.json",
        (json.dumps(metadata, indent=2, sort_keys=False) + "\n").encode("utf-8"),
    )
    return output_dir


def load_validated_artifact_metadata(
    artifact_dir: Path,
    source_repo: Path,
    *,
    require_resident_runtime: bool = False,
) -> dict:
    artifact_dir = artifact_dir.resolve()
    source_repo = source_repo.resolve()
    metadata_path = artifact_dir / "metadata.json"
    artifact_path = artifact_dir / "packed_transformer.safetensors"
    if not metadata_path.is_file() or not artifact_path.is_file():
        fail(f"artifact dir missing metadata or safetensors: {artifact_dir}")
    try:
        metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        fail(f"invalid artifact metadata {metadata_path}: {exc}")
    if metadata.get("format_version") != FORMAT_VERSION:
        fail(f"unsupported format_version: {metadata.get('format_version')}")
    if metadata.get("artifact_kind") != ARTIFACT_KIND:
        fail(f"unexpected artifact_kind: {metadata.get('artifact_kind')}")

    config_path = source_repo / "transformer" / "config.json"
    checkpoint_path = source_transformer_checkpoint(source_repo)
    expected_config_hash = sha256_file(config_path)
    expected_checkpoint_hash = sha256_file(checkpoint_path)
    try:
        source_metadata = metadata["source"]
        recorded_config_hash = source_metadata["transformer_config_sha256"]
        recorded_checkpoint_hash = source_metadata[
            "transformer_checkpoint_sha256"
        ]
        model_metadata = metadata["model"]
        recorded_depth = int(model_metadata["depth"])
        recorded_targets = metadata["targets"]
        recorded_non_target_keys = metadata["non_target_keys"]
    except (KeyError, TypeError, ValueError) as exc:
        fail(f"artifact metadata schema is incomplete or invalid: {exc}")
    if not isinstance(recorded_targets, list):
        fail("artifact metadata targets must be a list")
    if not isinstance(recorded_non_target_keys, list) or not all(
        isinstance(key, str) for key in recorded_non_target_keys
    ):
        fail("artifact metadata non_target_keys must be a list of strings")
    if recorded_config_hash != expected_config_hash:
        fail("transformer config hash mismatch")
    if recorded_checkpoint_hash != expected_checkpoint_hash:
        fail("transformer checkpoint hash mismatch")

    config = load_transformer_config(source_repo)
    if recorded_depth != int(config["depth"]):
        fail("artifact model depth does not match source config")
    specs = build_target_specs(int(config["depth"]))
    expected_modules = [spec.module_key for spec in specs]
    if not all(isinstance(entry, dict) for entry in recorded_targets):
        fail("artifact metadata target entries must be objects")
    recorded_modules = [entry.get("module_key") for entry in recorded_targets]
    if recorded_modules != expected_modules:
        fail("artifact target allowlist/order mismatch")
    target_source_keys = {
        key for spec in specs for key in (spec.weight_key, spec.bias_key)
    }
    expected_artifact_keys = {
        key
        for spec in specs
        for key in (
            spec.artifact_weight_key,
            spec.artifact_scale_key,
            spec.artifact_tensor_scale_key,
            spec.artifact_bias_key,
        )
    }
    with safe_open(artifact_path, framework="pt", device="cpu") as artifact_handle:
        actual_artifact_keys = set(artifact_handle.keys())
        header_metadata = artifact_handle.metadata()
    if actual_artifact_keys != expected_artifact_keys:
        fail("artifact tensor key coverage mismatch")
    if header_metadata not in (
        CANONICAL_SAFETENSORS_METADATA,
        LEGACY_SAFETENSORS_METADATA,
    ):
        fail("artifact safetensors header metadata mismatch")

    with safe_open(checkpoint_path, framework="pt", device="cpu") as source_handle:
        source_keys = set(source_handle.keys())
    missing_target_keys = sorted(target_source_keys - source_keys)
    if missing_target_keys:
        fail(
            "source checkpoint is missing target tensors, first="
            f"{missing_target_keys[0]}"
        )
    expected_non_target_keys = sorted(source_keys - target_source_keys)
    if recorded_non_target_keys != expected_non_target_keys:
        fail("artifact non-target source manifest mismatch")

    if require_resident_runtime:
        hashes = metadata.get("library_hashes", {})
        if not isinstance(hashes, dict):
            fail("artifact metadata library_hashes must be an object")
        if hashes.get("resident_source_sha256") != sha256_file(RESIDENT_SOURCE):
            fail("resident source hash mismatch")
        if hashes.get("resident_library_sha256") != sha256_file(RESIDENT_LIBRARY):
            fail("resident library hash mismatch")
    return metadata


def validate_artifact(artifact_dir: Path, source_repo: Path) -> None:
    artifact_dir = artifact_dir.resolve()
    source_repo = source_repo.resolve()
    metadata = load_validated_artifact_metadata(artifact_dir, source_repo)
    artifact_path = artifact_dir / "packed_transformer.safetensors"
    checkpoint_path = source_transformer_checkpoint(source_repo)
    config = load_transformer_config(source_repo)
    specs = {
        spec.module_key: spec for spec in build_target_specs(int(config["depth"]))
    }

    with safe_open(artifact_path, framework="pt", device="cpu") as artifact_handle, safe_open(
        checkpoint_path, framework="pt", device="cpu"
    ) as source_handle:
        for entry in metadata["targets"]:
            spec = specs[entry["module_key"]]
            weight = source_handle.get_tensor(spec.weight_key)
            bias = source_handle.get_tensor(spec.bias_key)
            packed_weight, packed_scales, weight_tensor_scale, global_amax = pack_weight_tensor(weight)

            candidate_weight = artifact_handle.get_tensor(spec.artifact_weight_key)
            candidate_scales = artifact_handle.get_tensor(spec.artifact_scale_key)
            candidate_tensor_scale = artifact_handle.get_tensor(spec.artifact_tensor_scale_key)
            candidate_bias = artifact_handle.get_tensor(spec.artifact_bias_key)

            if not torch.equal(candidate_weight, packed_weight):
                fail(f"packed weight mismatch for {spec.module_key}")
            if not torch.equal(candidate_scales, packed_scales):
                fail(f"packed scales mismatch for {spec.module_key}")
            if not torch.equal(candidate_tensor_scale, weight_tensor_scale):
                fail(f"tensor scale mismatch for {spec.module_key}")
            if not torch.equal(candidate_bias, bias):
                fail(f"bias mismatch for {spec.module_key}")
            if abs(float(entry["weight_global_amax"]) - global_amax) > 0.0:
                fail(f"global amax mismatch for {spec.module_key}")


def load_clean_transformer_from_artifact(
    artifact_dir: Path,
    source_repo: Path,
    assign_non_target: bool = True,
) -> nn.Module:
    artifact_dir = artifact_dir.resolve()
    source_repo = source_repo.resolve()
    metadata = load_validated_artifact_metadata(artifact_dir, source_repo)
    target_specs = build_target_specs(int(metadata["model"]["depth"]))
    model = instantiate_mage_transformer_on_meta(source_repo)
    replace_targets_with_artifact_modules(model, artifact_dir / "packed_transformer.safetensors", target_specs)

    if assign_non_target:
        skip_keys = {
            key
            for spec in target_specs
            for key in (spec.weight_key, spec.bias_key)
        }
        source_tensor_keys_read: list[str] = []
        with safe_open(source_transformer_checkpoint(source_repo), framework="pt", device="cpu") as source_handle:
            for key in source_handle.keys():
                if key in skip_keys:
                    continue
                tensor = source_handle.get_tensor(key)
                source_tensor_keys_read.append(key)
                assign_tensor_by_name(model, key, tensor)
        target_reads = sorted(set(source_tensor_keys_read) & skip_keys)
        if target_reads:
            fail(f"clean CPU loader read target source tensors: {target_reads[0]}")
        meta_parameters = [
            name for name, parameter in model.named_parameters() if parameter.is_meta
        ]
        meta_buffers = [
            name for name, buffer in model.named_buffers() if buffer.is_meta
        ]
        if meta_parameters or meta_buffers:
            first = (meta_parameters + meta_buffers)[0]
            fail(f"clean CPU loader left unresolved meta tensors, first={first}")
        materialize_mage_rope_tensor_attributes(model)
    return model


def load_clean_resident_transformer(
    artifact_dir: Path,
    source_repo: Path,
    device: torch.device,
) -> tuple[nn.Module, dict]:
    return _load_clean_cuda_transformer(
        artifact_dir,
        source_repo,
        device,
        replace_targets_with_resident_modules,
    )


def load_clean_native_resident_transformer(
    artifact_dir: Path,
    source_repo: Path,
    device: torch.device,
) -> tuple[nn.Module, dict]:
    return _load_clean_cuda_transformer(
        artifact_dir,
        source_repo,
        device,
        replace_targets_with_native_resident_modules,
    )


def _load_clean_cuda_transformer(
    artifact_dir: Path,
    source_repo: Path,
    device: torch.device,
    target_replacement_fn: Callable[[nn.Module, Path, Iterable[TargetSpec], torch.device], None],
) -> tuple[nn.Module, dict]:
    artifact_dir = artifact_dir.resolve()
    source_repo = source_repo.resolve()
    metadata = load_validated_artifact_metadata(
        artifact_dir,
        source_repo,
        require_resident_runtime=True,
    )
    target_specs = build_target_specs(int(metadata["model"]["depth"]))
    target_source_keys = {
        key
        for spec in target_specs
        for key in (spec.weight_key, spec.bias_key)
    }
    model = instantiate_mage_transformer_on_meta(source_repo)
    target_replacement_fn(
        model,
        artifact_dir / "packed_transformer.safetensors",
        target_specs,
        device,
    )

    loaded_source_keys: list[str] = []
    skipped_target_source_keys: list[str] = []
    source_tensor_keys_read: list[str] = []
    with safe_open(
        source_transformer_checkpoint(source_repo),
        framework="pt",
        device="cpu",
    ) as source_handle:
        source_keys = list(source_handle.keys())
        missing_target_keys = sorted(target_source_keys - set(source_keys))
        if missing_target_keys:
            fail(
                "source checkpoint is missing target tensors, first="
                f"{missing_target_keys[0]}"
            )
        for key in source_keys:
            if key in target_source_keys:
                skipped_target_source_keys.append(key)
                continue
            tensor = source_handle.get_tensor(key)
            source_tensor_keys_read.append(key)
            assign_tensor_by_name(model, key, tensor.to(device))
            loaded_source_keys.append(key)

    materialized_tensor_attributes = materialize_mage_rope_tensor_attributes(model)
    target_source_reads = sorted(
        set(source_tensor_keys_read) & target_source_keys
    )
    meta_parameters = [
        name for name, parameter in model.named_parameters() if parameter.is_meta
    ]
    meta_buffers = [
        name for name, buffer in model.named_buffers() if buffer.is_meta
    ]
    report = {
        "source_tensor_count_loaded": len(loaded_source_keys),
        "source_tensor_keys_loaded": loaded_source_keys,
        "source_tensor_count_read": len(source_tensor_keys_read),
        "source_tensor_keys_read": source_tensor_keys_read,
        "target_source_tensor_count_skipped": len(skipped_target_source_keys),
        "target_source_tensor_keys_skipped": skipped_target_source_keys,
        "target_source_tensor_reads": len(target_source_reads),
        "target_source_tensor_keys_read": target_source_reads,
        "meta_parameter_names": meta_parameters,
        "meta_buffer_names": meta_buffers,
        "materialized_unregistered_tensor_attribute_names": (
            materialized_tensor_attributes
        ),
        "unregistered_meta_tensor_attribute_names": (
            unregistered_meta_tensor_attribute_names(model)
        ),
    }
    return model.eval().requires_grad_(False), report


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)

    pack_parser = subparsers.add_parser("pack", help="build a packed artifact directory")
    pack_parser.add_argument("--source-repo", type=Path, required=True)
    pack_parser.add_argument("--output-dir", type=Path, required=True)

    validate_parser = subparsers.add_parser("validate", help="recompute and validate a packed artifact")
    validate_parser.add_argument("--artifact-dir", type=Path, required=True)
    validate_parser.add_argument("--source-repo", type=Path, required=True)

    plan_load_parser = subparsers.add_parser("plan-load", help="instantiate on meta and replace target modules")
    plan_load_parser.add_argument("--artifact-dir", type=Path, required=True)
    plan_load_parser.add_argument("--source-repo", type=Path, required=True)
    plan_load_parser.add_argument("--skip-non-target", action="store_true")

    runtime_parser = subparsers.add_parser(
        "validate-runtime",
        help="placeholder for future single-GPU resident validation",
    )
    runtime_parser.add_argument("--artifact-dir", type=Path, required=True)
    runtime_parser.add_argument("--source-repo", type=Path, required=True)
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    try:
        if args.command == "pack":
            artifact_dir = pack_artifact(args.source_repo, args.output_dir)
            print(artifact_dir)
            return 0
        if args.command == "validate":
            validate_artifact(args.artifact_dir, args.source_repo)
            print("ok")
            return 0
        if args.command == "plan-load":
            model = load_clean_transformer_from_artifact(
                args.artifact_dir, args.source_repo, assign_non_target=not args.skip_non_target
            )
            packed_count = sum(
                1 for _name, module in model.named_modules() if isinstance(module, PackedNvfp4LinearArtifactModule)
            )
            meta_parameters = [
                name for name, parameter in model.named_parameters() if parameter.is_meta
            ]
            meta_buffers = [
                name for name, buffer in model.named_buffers() if buffer.is_meta
            ]
            print(
                json.dumps(
                    {
                        "packed_module_count": packed_count,
                        "meta_parameter_count": len(meta_parameters),
                        "meta_buffer_count": len(meta_buffers),
                        "non_target_assignment_skipped": bool(args.skip_non_target),
                    },
                    indent=2,
                )
            )
            return 0
        if args.command == "validate-runtime":
            fail(
                "validate-runtime is intentionally not implemented in this CPU-only slice. "
                "Use the future CUDA resident path on CUDA_VISIBLE_DEVICES=3."
            )
        fail(f"unsupported command: {args.command}")
    except PackedArtifactError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())