multimodalart/minimax-h3-reference

1"""Bind AoTI constants that `torch.export` lifted anonymously.2 3`spaces.zero.torch.aoti.LazyAOTIModel` binds a package's constants by intersecting the module's `state_dict()` with4`compiled_model.get_constant_fqns()`, and keeps whatever it cannot match. `torch.export` only gives a lifted tensor a5real FQN when it was a registered parameter or buffer; anything else is named `_tensor_constant<N>`, which no6`state_dict()` can contain, so the compiled model runs against constants nobody set — a SIGSEGV rather than an error.7 8`write_constant_aliases` records the real names on the compile side; `apply_spaces_constant_binding_patch` uses that9sidecar on the load side, falls back to matching by dtype+shape, and raises if the binding is still not total.10"""11 12from __future__ import annotations13 14import io15import json16import re17import zipfile18from pathlib import Path19 20import torch21 22ALIASES_FILENAME = "constant_aliases.json"23 24_DTYPES = {25    "float32": torch.float32, "float64": torch.float64, "float16": torch.float16,26    "bfloat16": torch.bfloat16, "float8_e4m3fn": torch.float8_e4m3fn,27    "float8_e5m2": torch.float8_e5m2, "float8_e4m3fnuz": torch.float8_e4m3fnuz,28    "float8_e5m2fnuz": torch.float8_e5m2fnuz, "int8": torch.int8, "uint8": torch.uint8,29    "int16": torch.int16, "int32": torch.int32, "int64": torch.int64, "bool": torch.bool,30}31 32 33# --------------------------------------------------------------------------- compile side34 35 36def register_loose_tensors(module: torch.nn.Module, prefix: str = "") -> list[str]:37    """Re-register plain tensor attributes as buffers so `torch.export` gives them real FQNs.38 39    Run on the shallow clone, right before `torch.export.export`. Returns the names it re-registered.40    """41    registered = []42    for name, value in list(vars(module).items()):43        if not isinstance(value, torch.Tensor) or name.startswith("_"):44            continue45        if name in module._parameters or name in module._buffers:46            continue47        object.__delattr__(module, name)48        module.register_buffer(name, value, persistent=True)49        registered.append(f"{prefix}{name}")50    for child_name, child in module.named_children():51        registered += register_loose_tensors(child, f"{prefix}{child_name}.")52    return registered53 54 55def constant_aliases_from_exported_program(exported_program) -> dict[str, str]:56    """`{'_tensor_constant<N>': '<real dotted fqn>'}` for every anonymously lifted constant.57 58    AOT Inductor numbers its slots in the order the `CONSTANT_TENSOR` inputs appear in the graph signature, which59    still carries each one's real FQN.60    """61    targets = [62        spec.target63        for spec in exported_program.graph_signature.input_specs64        if spec.kind.name == "CONSTANT_TENSOR"65    ]66    return {f"_tensor_constant{index}": target for index, target in enumerate(targets)}67 68 69def write_constant_aliases(package_dir, exported_program, submodule: str | None = None) -> Path | None:70    """Drop the alias sidecar next to the `package.pt2` `aoti_compile_and_save` just wrote."""71    aliases = constant_aliases_from_exported_program(exported_program)72    if not aliases:73        return None74    subdir = Path(package_dir) / ("submodules/" + submodule if submodule else "root")75    path = subdir / ALIASES_FILENAME76    path.write_text(json.dumps(aliases, indent=2))77    return path78 79 80# --------------------------------------------------------------------------- load side81 82 83def _package_constants_info(archive_file) -> list[dict]:84    """Read `constants_info_` (dtype, shape, in slot order) out of a `.pt2`'s wrapper source."""85    if isinstance(archive_file, (str, Path)):86        handle: object = str(archive_file)87    else:88        position = archive_file.tell()89        archive_file.seek(0)90        handle = io.BytesIO(archive_file.read())91        archive_file.seek(position)92    with zipfile.ZipFile(handle) as archive:  # pyright: ignore[reportArgumentType]93        names = [n for n in archive.namelist() if n.endswith(".wrapper.cpp")]94        if not names:95            return []96        source = archive.read(names[0]).decode()97    info: dict[int, dict] = {}98    for match in re.finditer(r"constants_info_\[(\d+)\]\.(\w+) = ([^;]+);", source):99        index, field, value = int(match.group(1)), match.group(2), match.group(3).strip()100        entry = info.setdefault(index, {})101        if field == "dtype":102            entry["dtype"] = _DTYPES.get(value.replace("cached_torch_dtype_", ""))103        elif field == "shape":104            entry["shape"] = tuple(int(x) for x in re.findall(r"-?\d+", value))105        elif field in ("name", "original_fqn"):106            entry[field] = value.strip('"')107    return [info[index] for index in sorted(info)]108 109 110def resolve_constant_map(111    archive_file,112    constant_fqns,113    weights: dict[str, torch.Tensor],114    aliases=None,115    allow_shape_fallback: bool = False,116):117    """Map every compiled constant FQN onto one of `weights`, or report what is left over."""118    constant_map = {name: weights[name] for name in constant_fqns if name in weights}119    missing = [name for name in constant_fqns if name not in constant_map]120    if not missing:121        return constant_map, []122 123    aliases = aliases or {}124    for name in list(missing):125        target = aliases.get(name)126        if target is not None and target in weights:127            constant_map[name] = weights[target]128            missing.remove(name)129    if not missing or not allow_shape_fallback:130        return constant_map, missing131 132    # Match by dtype+shape against the unclaimed `state_dict()` entries, preserving each side's own order inside a133    # (dtype, shape) group. `get_constant_fqns()` returns slots lexicographically (`_tensor_constant10` before134    # `_tensor_constant2`), so the package's own `constants_info_` index is the only correct order to walk them in.135    info = _package_constants_info(archive_file)136    by_name = {entry.get("name"): entry for entry in info}137    slot_index = {entry.get("name"): index for index, entry in enumerate(info)}138    taken = {id(tensor) for tensor in constant_map.values()}139    buckets: dict[tuple, list[torch.Tensor]] = {}140    for tensor in weights.values():141        if id(tensor) not in taken:142            buckets.setdefault((tensor.dtype, tuple(tensor.shape)), []).append(tensor)143    for name in sorted(list(missing), key=lambda n: slot_index.get(n, 1 << 30)):144        entry = by_name.get(name)145        if entry is None or entry.get("dtype") is None:146            continue147        bucket = buckets.get((entry["dtype"], entry["shape"]))148        if bucket:149            constant_map[name] = bucket.pop(0)150            missing.remove(name)151    return constant_map, missing152 153 154def apply_spaces_constant_binding_patch(strict: bool = True, allow_shape_fallback: bool = False):155    """Make `spaces`' AoTI loader bind anonymous constants, and fail loudly if it still cannot.156 157    Call once, before any `spaces.aoti_*` loading. Idempotent.158    """159    from spaces.zero.torch import aoti as spaces_aoti160 161    if getattr(spaces_aoti.LazyAOTIModel, "_constant_binding_patched", False):162        return163 164    original_call = spaces_aoti.LazyAOTIModel.__call__165 166    def patched_call(self, weights, check_full_update, *args, **kwargs):167        compiled_model = self.compiled_model.get()168        if compiled_model is None:169            with spaces_aoti._register_aoti_cleanup():170                compiled_model = torch._inductor.aoti_load_package(self.archive_file)171            self.compiled_model.set(compiled_model)172        loaded = self.loaded_weights.get()173        if loaded is None or loaded is not weights:174            fqns = compiled_model.get_constant_fqns()175            aliases = getattr(self, "_constant_aliases", None)176            if aliases is None:177                aliases = {}178                if isinstance(self.archive_file, (str, Path)):179                    sidecar = Path(self.archive_file).with_name(ALIASES_FILENAME)180                    if sidecar.is_file():181                        aliases = json.loads(sidecar.read_text())182                self._constant_aliases = aliases183            constant_map, missing = resolve_constant_map(184                self.archive_file, fqns, weights, aliases, allow_shape_fallback185            )186            if missing and strict:187                raise RuntimeError(188                    f"{len(missing)} of {len(fqns)} AoTI constants could not be bound to the module's "189                    f"state_dict: {missing[:8]}. Anonymous `_tensor_constant*` names mean the export saw "190                    f"plain tensor attributes rather than registered parameters or buffers. Register them "191                    f"(or write a {ALIASES_FILENAME} sidecar at compile time) — binding them partially "192                    f"would leave the compiled model dereferencing unset constants."193                )194            compiled_model.load_constants(195                constant_map, check_full_update=check_full_update and not missing, user_managed=True196            )197            self.loaded_weights.set(weights)198        return compiled_model(*args, **kwargs)199 200    spaces_aoti.LazyAOTIModel.__call__ = patched_call201    spaces_aoti.LazyAOTIModel._constant_binding_patched = True202    spaces_aoti.LazyAOTIModel._original_call = original_call203