Spaces:
Running on Zero
Running on Zero
File size: 3,180 Bytes
ee22178 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """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)
|