Upload spaces_constant_binding_patch.py with huggingface_hub

spaces_constant_binding_patch.py

@@ -0,0 +1,202 @@
+"""Bind AoTI constants that `torch.export` lifted anonymously.
+
+`spaces.zero.torch.aoti.LazyAOTIModel` binds a package's constants by intersecting the module's `state_dict()` with
+`compiled_model.get_constant_fqns()`, and keeps whatever it cannot match. `torch.export` only gives a lifted tensor a
+real FQN when it was a registered parameter or buffer; anything else is named `_tensor_constant<N>`, which no
+`state_dict()` can contain, so the compiled model runs against constants nobody set — a SIGSEGV rather than an error.
+
+`write_constant_aliases` records the real names on the compile side; `apply_spaces_constant_binding_patch` uses that
+sidecar on the load side, falls back to matching by dtype+shape, and raises if the binding is still not total.
+"""
+
+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.
+
+    Run on the shallow clone, right before `torch.export.export`. Returns the names it re-registered.
+    """
+    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 slots in the order the `CONSTANT_TENSOR` inputs appear in the graph signature, which
+    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 report what is left over."""
+    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, []
+
+    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
+
+    # Match by dtype+shape against the unclaimed `state_dict()` entries, preserving each side's own order inside a
+    # (dtype, shape) group. `get_constant_fqns()` returns slots lexicographically (`_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