#!/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())