aoti.py

"""Compile repeated transformer blocks once on the Space's actual ZeroGPU device."""

import logging
from pathlib import Path
import time

import spaces
import torch
from spaces.zero.torch.aoti import LazyAOTIModel

logger = logging.getLogger(__name__)
BLOCK_GROUPS = ("text_embedder.blocks", "layers")


def export_block(block, args):
    # The transformer evaluates each sample separately at its true sequence length.
    sequence = torch.export.Dim("sequence", min=1)
    shapes = ({1: sequence}, {1: sequence}) + ({},) * (len(args) - 2)
    return torch.export.export(block, args=args, dynamic_shapes=shapes)


def prepare_aoti(pipelines, package_dir=".cache/aoti"):
    pipelines = tuple(pipelines)
    transformer = pipelines[0].transformer
    package_dir = Path(package_dir)

    @spaces.GPU(duration=1500)
    @torch.no_grad()
    def compile_blocks():
        started = time.monotonic()
        logger.info("AOT compilation starting on %s", torch.cuda.get_device_name())
        cfg = transformer.config
        dtype = next(transformer.parameters()).dtype
        # Capture real input layouts/dtypes, without encoding a prompt or decoding an image.
        inputs = dict(
            hidden_states=torch.zeros(1, cfg.latent_channels, 64, 64, device="cuda"),
            timestep=torch.tensor([500.0], device="cuda"),
            encoder_hidden_states=torch.zeros(1, 32, cfg.text_dim, device="cuda", dtype=dtype),
        )
        for name in BLOCK_GROUPS:
            block = transformer.get_submodule(name)[0]
            with spaces.aoti_capture(block) as call:
                transformer(**inputs)
            exported = export_block(block, call.args)
            spaces.aoti_compile_and_save(
                package_dir, exported, submodule=name,
                inductor_configs={"dynamic_scale_rblock": False},
            )
            # Verify that the shared graph uses each model's own weights.
            compiled = LazyAOTIModel(package_dir / "submodules" / name / "package.pt2")
            for pipe in pipelines:
                candidate = pipe.transformer.get_submodule(name)[-1]
                expected = candidate(*call.args).float()
                actual = compiled.with_weights(candidate.state_dict())(*call.args).float()
                relative_error = ((actual - expected).square().mean().sqrt()
                                  / expected.square().mean().sqrt().clamp_min(1e-6)).item()
                if not torch.isfinite(actual).all() or relative_error > 0.02:
                    raise RuntimeError(f"AOT validation failed for {name}: relative RMS error {relative_error}")
                logger.info("AOT validated %s revision=%s relative_rms_error=%.6f",
                            name, pipe._moss_revision, relative_error)
            logger.info("AOT compiled %s in %.1fs", name, time.monotonic() - started)
        return round(time.monotonic() - started, 3)

    elapsed = compile_blocks()
    for pipe in pipelines:
        spaces.aoti_load_from_package_dir(pipe.transformer, package_dir)
        pipe._execution_backend = "regional-aoti"
    logger.info("AOT ready for both models; startup compilation took %.1fs", elapsed)