Wire the AoTI load path (env-gated) and the sub-768p canvases

spaces_constant_binding_patch.py

@@ -0,0 +1,227 @@
+"""Bind AoTI constants that `torch.export` lifted anonymously.
+
+Problem
+-------
+`spaces.zero.torch.aoti.LazyAOTIModel` binds a compiled package's constants **by name**::
+
+    constant_fqns = compiled_model.get_constant_fqns()
+    constant_map = {name: tensor for name, tensor in weights.items() if name in constant_fqns}
+    compiled_model.load_constants(constant_map, check_full_update=check_full_update, user_managed=True)
+
+`torch.export` only gives a lifted tensor a real FQN when it was a registered parameter or buffer.
+Anything reached through a plain python attribute is classified `CONSTANT_TENSOR` and the compiled
+artifact names it `_tensor_constant<N>` — a name that can never appear in `state_dict()`. The
+intersection above is then empty, the dict comprehension silently drops every weight, and the
+compiled model runs against constants nobody ever set: a SIGSEGV rather than an error.
+
+This module fixes both halves:
+
+  * `write_constant_aliases(...)` — compile side. Records the exact
+    `_tensor_constant<N> -> real.dotted.fqn` mapping, which the `ExportedProgram` knows even when the
+    compiled package does not, into a `constant_aliases.json` sidecar next to `package.pt2`.
+
+  * `apply_spaces_constant_binding_patch()` — load side. Monkeypatches `LazyAOTIModel.__call__` so it
+    (1) uses that sidecar when present, (2) otherwise falls back to matching anonymous constants
+    against the leftover `state_dict()` entries by dtype+shape read out of the package's own
+    `wrapper.cpp`, and (3) **raises** if the binding is not total instead of segfaulting later.
+
+The load-side patch alone is enough to turn the crash into a clear diagnostic; with the sidecar it
+also makes the package work.
+"""
+
+from __future__ import annotations
+
+import io
+import json
+import re
+import zipfile
+from pathlib import Path
+
+import torch
+
+ALIASES_FILENAME = "constant_aliases.json"
+
+_DTYPES = {
+    "float32": torch.float32, "float64": torch.float64, "float16": torch.float16,
+    "bfloat16": torch.bfloat16, "float8_e4m3fn": torch.float8_e4m3fn,
+    "float8_e5m2": torch.float8_e5m2, "float8_e4m3fnuz": torch.float8_e4m3fnuz,
+    "float8_e5m2fnuz": torch.float8_e5m2fnuz, "int8": torch.int8, "uint8": torch.uint8,
+    "int16": torch.int16, "int32": torch.int32, "int64": torch.int64, "bool": torch.bool,
+}
+
+
+# --------------------------------------------------------------------------- compile side
+
+
+def register_loose_tensors(module: torch.nn.Module, prefix: str = "") -> list[str]:
+    """Re-register plain tensor attributes as buffers so `torch.export` gives them real FQNs.
+
+    Model-agnostic and numerics-preserving: it changes how a tensor is *registered*, never the tensor
+    and never the forward. Run it on the shallow clone right after
+    `unwrap_tensor_subclass_parameters`, immediately before `torch.export.export`. Returns the names
+    it re-registered, which is empty for a module that was already well-formed.
+    """
+    registered = []
+    for name, value in list(vars(module).items()):
+        if not isinstance(value, torch.Tensor) or name.startswith("_"):
+            continue
+        if name in module._parameters or name in module._buffers:
+            continue
+        object.__delattr__(module, name)
+        module.register_buffer(name, value, persistent=True)
+        registered.append(f"{prefix}{name}")
+    for child_name, child in module.named_children():
+        registered += register_loose_tensors(child, f"{prefix}{child_name}.")
+    return registered
+
+
+def constant_aliases_from_exported_program(exported_program) -> dict[str, str]:
+    """`{'_tensor_constant<N>': '<real dotted fqn>'}` for every anonymously lifted constant.
+
+    AOT Inductor numbers its `_tensor_constant<N>` slots in the order the `CONSTANT_TENSOR` inputs
+    appear in the export graph signature, and the signature still carries each one's real FQN.
+    """
+    targets = [
+        spec.target
+        for spec in exported_program.graph_signature.input_specs
+        if spec.kind.name == "CONSTANT_TENSOR"
+    ]
+    return {f"_tensor_constant{index}": target for index, target in enumerate(targets)}
+
+
+def write_constant_aliases(package_dir, exported_program, submodule: str | None = None) -> Path | None:
+    """Drop the alias sidecar next to the `package.pt2` `aoti_compile_and_save` just wrote."""
+    aliases = constant_aliases_from_exported_program(exported_program)
+    if not aliases:
+        return None
+    subdir = Path(package_dir) / ("submodules/" + submodule if submodule else "root")
+    path = subdir / ALIASES_FILENAME
+    path.write_text(json.dumps(aliases, indent=2))
+    return path
+
+
+# --------------------------------------------------------------------------- load side
+
+
+def _package_constants_info(archive_file) -> list[dict]:
+    """Read `constants_info_` (dtype, shape, in slot order) out of a `.pt2`'s wrapper source."""
+    if isinstance(archive_file, (str, Path)):
+        handle: object = str(archive_file)
+    else:
+        position = archive_file.tell()
+        archive_file.seek(0)
+        handle = io.BytesIO(archive_file.read())
+        archive_file.seek(position)
+    with zipfile.ZipFile(handle) as archive:  # pyright: ignore[reportArgumentType]
+        names = [n for n in archive.namelist() if n.endswith(".wrapper.cpp")]
+        if not names:
+            return []
+        source = archive.read(names[0]).decode()
+    info: dict[int, dict] = {}
+    for match in re.finditer(r"constants_info_\[(\d+)\]\.(\w+) = ([^;]+);", source):
+        index, field, value = int(match.group(1)), match.group(2), match.group(3).strip()
+        entry = info.setdefault(index, {})
+        if field == "dtype":
+            entry["dtype"] = _DTYPES.get(value.replace("cached_torch_dtype_", ""))
+        elif field == "shape":
+            entry["shape"] = tuple(int(x) for x in re.findall(r"-?\d+", value))
+        elif field in ("name", "original_fqn"):
+            entry[field] = value.strip('"')
+    return [info[index] for index in sorted(info)]
+
+
+def resolve_constant_map(
+    archive_file,
+    constant_fqns,
+    weights: dict[str, torch.Tensor],
+    aliases=None,
+    allow_shape_fallback: bool = False,
+):
+    """Map every compiled constant FQN onto one of `weights`, or explain why it cannot."""
+    constant_map = {name: weights[name] for name in constant_fqns if name in weights}
+    missing = [name for name in constant_fqns if name not in constant_map]
+    if not missing:
+        return constant_map, []
+
+    # 1. the exact mapping, if the compile side recorded one
+    aliases = aliases or {}
+    for name in list(missing):
+        target = aliases.get(name)
+        if target is not None and target in weights:
+            constant_map[name] = weights[target]
+            missing.remove(name)
+    if not missing or not allow_shape_fallback:
+        return constant_map, missing
+
+    # 2. otherwise match by dtype+shape against the state_dict entries nobody claimed, preserving
+    #    each side's own order inside a (dtype, shape) group. `get_constant_fqns()` returns the
+    #    slots in *lexicographic* order (`_tensor_constant10` before `_tensor_constant2`), so the
+    #    package's own `constants_info_` index is the only correct order to walk them in.
+    info = _package_constants_info(archive_file)
+    by_name = {entry.get("name"): entry for entry in info}
+    slot_index = {entry.get("name"): index for index, entry in enumerate(info)}
+    taken = {id(tensor) for tensor in constant_map.values()}
+    buckets: dict[tuple, list[torch.Tensor]] = {}
+    for tensor in weights.values():
+        if id(tensor) not in taken:
+            buckets.setdefault((tensor.dtype, tuple(tensor.shape)), []).append(tensor)
+    for name in sorted(list(missing), key=lambda n: slot_index.get(n, 1 << 30)):
+        entry = by_name.get(name)
+        if entry is None or entry.get("dtype") is None:
+            continue
+        bucket = buckets.get((entry["dtype"], entry["shape"]))
+        if bucket:
+            constant_map[name] = bucket.pop(0)
+            missing.remove(name)
+    return constant_map, missing
+
+
+def apply_spaces_constant_binding_patch(strict: bool = True, allow_shape_fallback: bool = False):
+    """Make `spaces`' AoTI loader bind anonymous constants, and fail loudly if it still cannot.
+
+    Call once, before any `spaces.aoti_*` loading. Idempotent.
+    """
+    from spaces.zero.torch import aoti as spaces_aoti
+
+    if getattr(spaces_aoti.LazyAOTIModel, "_constant_binding_patched", False):
+        return
+
+    original_call = spaces_aoti.LazyAOTIModel.__call__
+
+    def patched_call(self, weights, check_full_update, *args, **kwargs):
+        compiled_model = self.compiled_model.get()
+        if compiled_model is None:
+            with spaces_aoti._register_aoti_cleanup():
+                compiled_model = torch._inductor.aoti_load_package(self.archive_file)
+            self.compiled_model.set(compiled_model)
+        loaded = self.loaded_weights.get()
+        if loaded is None or loaded is not weights:
+            fqns = compiled_model.get_constant_fqns()
+            aliases = getattr(self, "_constant_aliases", None)
+            if aliases is None:
+                aliases = {}
+                if isinstance(self.archive_file, (str, Path)):
+                    sidecar = Path(self.archive_file).with_name(ALIASES_FILENAME)
+                    if sidecar.is_file():
+                        aliases = json.loads(sidecar.read_text())
+                self._constant_aliases = aliases
+            constant_map, missing = resolve_constant_map(
+                self.archive_file, fqns, weights, aliases, allow_shape_fallback
+            )
+            if missing and strict:
+                raise RuntimeError(
+                    f"{len(missing)} of {len(fqns)} AoTI constants could not be bound to the module's "
+                    f"state_dict: {missing[:8]}. Anonymous `_tensor_constant*` names mean the export saw "
+                    f"plain tensor attributes rather than registered parameters or buffers. Register them "
+                    f"(or write a {ALIASES_FILENAME} sidecar at compile time) — binding them partially "
+                    f"would leave the compiled model dereferencing unset constants."
+                )
+            compiled_model.load_constants(
+                constant_map, check_full_update=check_full_update and not missing, user_managed=True
+            )
+            self.loaded_weights.set(weights)
+        return compiled_model(*args, **kwargs)
+
+    spaces_aoti.LazyAOTIModel.__call__ = patched_call
+    spaces_aoti.LazyAOTIModel._constant_binding_patched = True
+    spaces_aoti.LazyAOTIModel._original_call = original_call