runtime/standard_transformer.py

"""Load the complete sharded Mage-Flow NVFP4 transformer component."""

from __future__ import annotations

from contextlib import ExitStack
import json
from pathlib import Path
from typing import Any

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

from packed_artifact import (
    assign_tensor_by_name,
    build_target_specs,
    instantiate_mage_transformer_on_meta,
    materialize_mage_rope_tensor_attributes,
    set_child_module,
    unregistered_meta_tensor_attribute_names,
)
from torch_ops_native import (
    PackedNvfp4LinearNativeOp,
    initialize_native_sm120_op,
)


class StandardCheckpointError(RuntimeError):
    pass


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


def read_object(path: Path) -> dict[str, Any]:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        fail(f"cannot read JSON object {path}: {exc}")
    if not isinstance(value, dict):
        fail(f"expected a JSON object: {path}")
    return value


def _component_path(component_dir: Path, relative: str) -> Path:
    path = (component_dir / relative).resolve()
    if not path.is_relative_to(component_dir):
        fail(f"checkpoint index path escapes transformer component: {relative}")
    if not path.is_file():
        fail(f"checkpoint shard is missing: {relative}")
    return path


def _quantized_keys(module_key: str) -> dict[str, str]:
    return {
        "packed_weight": f"{module_key}.packed_weight",
        "weight_scales": f"{module_key}.weight_scales",
        "weight_scale": f"{module_key}.weight_scale",
        "bias": f"{module_key}.bias",
    }


def load_standard_native_transformer(
    repo_root: str | Path,
    device: torch.device,
) -> tuple[nn.Module, dict[str, Any]]:
    """Load a complete standard-layout component without a BF16 base download."""

    repo_root = Path(repo_root).resolve()
    component_dir = (repo_root / "transformer").resolve()
    if not component_dir.is_relative_to(repo_root) or not component_dir.is_dir():
        fail("repository has no transformer component")
    if device.type != "cuda":
        fail("the native resident transformer requires a CUDA destination")
    if not initialize_native_sm120_op(allow_python_schema_fallback=False):
        fail("compiled native SM120 torch op did not load")

    config = read_object(component_dir / "config.json")
    quant_config = config.get("quantization_config")
    if not isinstance(quant_config, dict):
        fail("transformer config has no quantization_config")
    if quant_config.get("quant_method") != "mage_flow_nvfp4":
        fail(
            "unexpected transformer quantization method: "
            f"{quant_config.get('quant_method')!r}"
        )
    if quant_config.get("quant_algo") != "NVFP4":
        fail("transformer config does not declare NVFP4")

    depth = int(config.get("depth", 0))
    target_specs = build_target_specs(depth)
    expected_targets = [spec.module_key for spec in target_specs]
    if quant_config.get("targets") != expected_targets:
        fail("transformer quantization target list is not canonical")

    metadata = read_object(component_dir / "nvfp4_metadata.json")
    if metadata.get("artifact_kind") != (
        "mage_flow_transformer_mlp_nvfp4_resident_v1"
    ):
        fail("unexpected transformer NVFP4 metadata kind")
    non_target_keys = metadata.get("non_target_keys")
    if not isinstance(non_target_keys, list) or not all(
        isinstance(key, str) for key in non_target_keys
    ):
        fail("transformer NVFP4 metadata has no non-target key list")

    index = read_object(
        component_dir / "diffusion_pytorch_model.safetensors.index.json"
    )
    weight_map = index.get("weight_map")
    if not isinstance(weight_map, dict) or not all(
        isinstance(key, str) and isinstance(value, str)
        for key, value in weight_map.items()
    ):
        fail("transformer checkpoint has no valid weight map")

    quantized_keys = {
        key
        for spec in target_specs
        for key in _quantized_keys(spec.module_key).values()
    }
    expected_keys = set(non_target_keys) | quantized_keys
    actual_keys = set(weight_map)
    if actual_keys != expected_keys:
        missing = sorted(expected_keys - actual_keys)
        unexpected = sorted(actual_keys - expected_keys)
        fail(
            "transformer checkpoint key coverage mismatch; "
            f"missing={missing[:1]}, unexpected={unexpected[:1]}"
        )
    original_target_weights = {spec.weight_key for spec in target_specs}
    leaked = sorted(actual_keys & original_target_weights)
    if leaked:
        fail(f"BF16 target weight leaked into quantized checkpoint: {leaked[0]}")

    shard_names = sorted(set(weight_map.values()))
    with ExitStack() as stack:
        handles = {
            name: stack.enter_context(
                safe_open(
                    _component_path(component_dir, name),
                    framework="pt",
                    device="cpu",
                )
            )
            for name in shard_names
        }

        def tensor(key: str) -> torch.Tensor:
            try:
                handle = handles[weight_map[key]]
            except KeyError:
                fail(f"tensor is absent from checkpoint index: {key}")
            if key not in handle.keys():
                fail(f"tensor is absent from its declared shard: {key}")
            return handle.get_tensor(key)

        model = instantiate_mage_transformer_on_meta(repo_root)
        for spec in target_specs:
            original = model.get_submodule(spec.module_key)
            if not isinstance(original, nn.Linear):
                fail(
                    f"expected target {spec.module_key} to be nn.Linear, "
                    f"found {type(original).__name__}"
                )
            keys = _quantized_keys(spec.module_key)
            replacement = PackedNvfp4LinearNativeOp(
                in_features=int(original.in_features),
                out_features=int(original.out_features),
                packed_weight=tensor(keys["packed_weight"]).to(device),
                weight_scales=tensor(keys["weight_scales"]).to(device),
                weight_scale=tensor(keys["weight_scale"]).to(device),
                bias=tensor(keys["bias"]).to(device),
            )
            set_child_module(model, spec.module_key, replacement)

        loaded_non_targets: list[str] = []
        for key in non_target_keys:
            assign_tensor_by_name(model, key, tensor(key).to(device))
            loaded_non_targets.append(key)

    materialized = materialize_mage_rope_tensor_attributes(model)
    meta_parameters = [
        name for name, value in model.named_parameters() if value.is_meta
    ]
    meta_buffers = [
        name for name, value in model.named_buffers() if value.is_meta
    ]
    unregistered_meta = unregistered_meta_tensor_attribute_names(model)
    if meta_parameters or meta_buffers or unregistered_meta:
        fail(
            "standard transformer loader left unresolved meta tensors: "
            f"{(meta_parameters + meta_buffers + unregistered_meta)[:4]}"
        )

    report = {
        "layout": "huggingface_sharded_component",
        "checkpoint_shard_count": len(shard_names),
        "checkpoint_tensor_count": len(actual_keys),
        "loaded_non_target_tensor_count": len(loaded_non_targets),
        "loaded_quantized_projection_count": len(target_specs),
        "bf16_target_weight_reads": 0,
        "meta_parameter_names": meta_parameters,
        "meta_buffer_names": meta_buffers,
        "materialized_unregistered_tensor_attribute_names": materialized,
    }
    return model.eval().requires_grad_(False), report


__all__ = [
    "StandardCheckpointError",
    "load_standard_native_transformer",
]