Add runtime/packed_artifact.py

runtime/packed_artifact.py

@@ -0,0 +1,1068 @@
+#!/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())