yn4989/minimax-h3

1"""Bind AoTI constants that `torch.export` lifted anonymously.2 3Problem4-------5`spaces.zero.torch.aoti.LazyAOTIModel` binds a compiled package's constants **by name**::6 7    constant_fqns = compiled_model.get_constant_fqns()8    constant_map = {name: tensor for name, tensor in weights.items() if name in constant_fqns}9    compiled_model.load_constants(constant_map, check_full_update=check_full_update, user_managed=True)10 11`torch.export` only gives a lifted tensor a real FQN when it was a registered parameter or buffer.12Anything reached through a plain python attribute is classified `CONSTANT_TENSOR` and the compiled13artifact names it `_tensor_constant<N>` — a name that can never appear in `state_dict()`. The14intersection above is then empty, the dict comprehension silently drops every weight, and the15compiled model runs against constants nobody ever set: a SIGSEGV rather than an error.16 17This module fixes both halves:18 19  * `write_constant_aliases(...)` — compile side. Records the exact20    `_tensor_constant<N> -> real.dotted.fqn` mapping, which the `ExportedProgram` knows even when the21    compiled package does not, into a `constant_aliases.json` sidecar next to `package.pt2`.22 23  * `apply_spaces_constant_binding_patch()` — load side. Monkeypatches `LazyAOTIModel.__call__` so it24    (1) uses that sidecar when present, (2) otherwise falls back to matching anonymous constants25    against the leftover `state_dict()` entries by dtype+shape read out of the package's own26    `wrapper.cpp`, and (3) **raises** if the binding is not total instead of segfaulting later.27 28The load-side patch alone is enough to turn the crash into a clear diagnostic; with the sidecar it29also makes the package work.30"""31 32from __future__ import annotations33 34import io35import json36import re37import zipfile38from pathlib import Path39 40import torch41 42ALIASES_FILENAME = "constant_aliases.json"43 44_DTYPES = {45    "float32": torch.float32, "float64": torch.float64, "float16": torch.float16,46    "bfloat16": torch.bfloat16, "float8_e4m3fn": torch.float8_e4m3fn,47    "float8_e5m2": torch.float8_e5m2, "float8_e4m3fnuz": torch.float8_e4m3fnuz,48    "float8_e5m2fnuz": torch.float8_e5m2fnuz, "int8": torch.int8, "uint8": torch.uint8,49    "int16": torch.int16, "int32": torch.int32, "int64": torch.int64, "bool": torch.bool,50}51 52 53# --------------------------------------------------------------------------- compile side54 55 56def register_loose_tensors(module: torch.nn.Module, prefix: str = "") -> list[str]:57    """Re-register plain tensor attributes as buffers so `torch.export` gives them real FQNs.58 59    Model-agnostic and numerics-preserving: it changes how a tensor is *registered*, never the tensor60    and never the forward. Run it on the shallow clone right after61    `unwrap_tensor_subclass_parameters`, immediately before `torch.export.export`. Returns the names62    it re-registered, which is empty for a module that was already well-formed.63    """64    registered = []65    for name, value in list(vars(module).items()):66        if not isinstance(value, torch.Tensor) or name.startswith("_"):67            continue68        if name in module._parameters or name in module._buffers:69            continue70        object.__delattr__(module, name)71        module.register_buffer(name, value, persistent=True)72        registered.append(f"{prefix}{name}")73    for child_name, child in module.named_children():74        registered += register_loose_tensors(child, f"{prefix}{child_name}.")75    return registered76 77 78def constant_aliases_from_exported_program(exported_program) -> dict[str, str]:79    """`{'_tensor_constant<N>': '<real dotted fqn>'}` for every anonymously lifted constant.80 81    AOT Inductor numbers its `_tensor_constant<N>` slots in the order the `CONSTANT_TENSOR` inputs82    appear in the export graph signature, and the signature still carries each one's real FQN.83    """84    targets = [85        spec.target86        for spec in exported_program.graph_signature.input_specs87        if spec.kind.name == "CONSTANT_TENSOR"88    ]89    return {f"_tensor_constant{index}": target for index, target in enumerate(targets)}90 91 92def write_constant_aliases(package_dir, exported_program, submodule: str | None = None) -> Path | None:93    """Drop the alias sidecar next to the `package.pt2` `aoti_compile_and_save` just wrote."""94    aliases = constant_aliases_from_exported_program(exported_program)95    if not aliases:96        return None97    subdir = Path(package_dir) / ("submodules/" + submodule if submodule else "root")98    path = subdir / ALIASES_FILENAME99    path.write_text(json.dumps(aliases, indent=2))100    return path101 102 103# --------------------------------------------------------------------------- load side104 105 106def _package_constants_info(archive_file) -> list[dict]:107    """Read `constants_info_` (dtype, shape, in slot order) out of a `.pt2`'s wrapper source."""108    if isinstance(archive_file, (str, Path)):109        handle: object = str(archive_file)110    else:111        position = archive_file.tell()112        archive_file.seek(0)113        handle = io.BytesIO(archive_file.read())114        archive_file.seek(position)115    with zipfile.ZipFile(handle) as archive:  # pyright: ignore[reportArgumentType]116        names = [n for n in archive.namelist() if n.endswith(".wrapper.cpp")]117        if not names:118            return []119        source = archive.read(names[0]).decode()120    info: dict[int, dict] = {}121    for match in re.finditer(r"constants_info_\[(\d+)\]\.(\w+) = ([^;]+);", source):122        index, field, value = int(match.group(1)), match.group(2), match.group(3).strip()123        entry = info.setdefault(index, {})124        if field == "dtype":125            entry["dtype"] = _DTYPES.get(value.replace("cached_torch_dtype_", ""))126        elif field == "shape":127            entry["shape"] = tuple(int(x) for x in re.findall(r"-?\d+", value))128        elif field in ("name", "original_fqn"):129            entry[field] = value.strip('"')130    return [info[index] for index in sorted(info)]131 132 133def resolve_constant_map(134    archive_file,135    constant_fqns,136    weights: dict[str, torch.Tensor],137    aliases=None,138    allow_shape_fallback: bool = False,139):140    """Map every compiled constant FQN onto one of `weights`, or explain why it cannot."""141    constant_map = {name: weights[name] for name in constant_fqns if name in weights}142    missing = [name for name in constant_fqns if name not in constant_map]143    if not missing:144        return constant_map, []145 146    # 1. the exact mapping, if the compile side recorded one147    aliases = aliases or {}148    for name in list(missing):149        target = aliases.get(name)150        if target is not None and target in weights:151            constant_map[name] = weights[target]152            missing.remove(name)153    if not missing or not allow_shape_fallback:154        return constant_map, missing155 156    # 2. otherwise match by dtype+shape against the state_dict entries nobody claimed, preserving157    #    each side's own order inside a (dtype, shape) group. `get_constant_fqns()` returns the158    #    slots in *lexicographic* order (`_tensor_constant10` before `_tensor_constant2`), so the159    #    package's own `constants_info_` index is the only correct order to walk them in.160    info = _package_constants_info(archive_file)161    by_name = {entry.get("name"): entry for entry in info}162    slot_index = {entry.get("name"): index for index, entry in enumerate(info)}163    taken = {id(tensor) for tensor in constant_map.values()}164    buckets: dict[tuple, list[torch.Tensor]] = {}165    for tensor in weights.values():166        if id(tensor) not in taken:167            buckets.setdefault((tensor.dtype, tuple(tensor.shape)), []).append(tensor)168    for name in sorted(list(missing), key=lambda n: slot_index.get(n, 1 << 30)):169        entry = by_name.get(name)170        if entry is None or entry.get("dtype") is None:171            continue172        bucket = buckets.get((entry["dtype"], entry["shape"]))173        if bucket:174            constant_map[name] = bucket.pop(0)175            missing.remove(name)176    return constant_map, missing177 178 179def apply_spaces_constant_binding_patch(strict: bool = True, allow_shape_fallback: bool = False):180    """Make `spaces`' AoTI loader bind anonymous constants, and fail loudly if it still cannot.181 182    Call once, before any `spaces.aoti_*` loading. Idempotent.183    """184    from spaces.zero.torch import aoti as spaces_aoti185 186    if getattr(spaces_aoti.LazyAOTIModel, "_constant_binding_patched", False):187        return188 189    original_call = spaces_aoti.LazyAOTIModel.__call__190 191    def patched_call(self, weights, check_full_update, *args, **kwargs):192        compiled_model = self.compiled_model.get()193        if compiled_model is None:194            with spaces_aoti._register_aoti_cleanup():195                compiled_model = torch._inductor.aoti_load_package(self.archive_file)196            self.compiled_model.set(compiled_model)197        loaded = self.loaded_weights.get()198        if loaded is None or loaded is not weights:199            fqns = compiled_model.get_constant_fqns()200            aliases = getattr(self, "_constant_aliases", None)201            if aliases is None:202                aliases = {}203                if isinstance(self.archive_file, (str, Path)):204                    sidecar = Path(self.archive_file).with_name(ALIASES_FILENAME)205                    if sidecar.is_file():206                        aliases = json.loads(sidecar.read_text())207                self._constant_aliases = aliases208            constant_map, missing = resolve_constant_map(209                self.archive_file, fqns, weights, aliases, allow_shape_fallback210            )211            if missing and strict:212                raise RuntimeError(213                    f"{len(missing)} of {len(fqns)} AoTI constants could not be bound to the module's "214                    f"state_dict: {missing[:8]}. Anonymous `_tensor_constant*` names mean the export saw "215                    f"plain tensor attributes rather than registered parameters or buffers. Register them "216                    f"(or write a {ALIASES_FILENAME} sidecar at compile time) — binding them partially "217                    f"would leave the compiled model dereferencing unset constants."218                )219            compiled_model.load_constants(220                constant_map, check_full_update=check_full_update and not missing, user_managed=True221            )222            self.loaded_weights.set(weights)223        return compiled_model(*args, **kwargs)224 225    spaces_aoti.LazyAOTIModel.__call__ = patched_call226    spaces_aoti.LazyAOTIModel._constant_binding_patched = True227    spaces_aoti.LazyAOTIModel._original_call = original_call228