diff --git a/configs/z_image_turbo.yaml b/configs/z_image_turbo.yaml index b2bf0f2fa2..228890d047 100644 --- a/configs/z_image_turbo.yaml +++ b/configs/z_image_turbo.yaml @@ -18,26 +18,29 @@ model_kwargs: # [norms] for reference-faithful output. compile_exact_ops: false accelerator_graph: true - capture_sizes: [[1024, 1024]] + # [height, width]: 1024^2 and 1344x768, the two sizes of the image benchmark set; an uncaptured size compiles per + # batch size inside its first requests and decodes on the eager VAE. + capture_sizes: [[1024, 1024], [768, 1344]] # Captions pad to multiples of 32 up to max_seq_len, so cap_len has 16 possible # values -- and it is an INPUT, not a shape: the DiT substitutes a learned # cap_pad_token at every padded position and attends over it, so a prompt cannot # be padded up into a wider bucket without changing the image. A cap_len with no # captured graph falls back to the eager batched path. # - # The capture budget is len(capture_caption_lengths) x len(capture_batch_sizes) graphs, - # 16 here. A step of 3 or 5-7 rows replays the next captured size with dummy rows. - # Capturing every batch size up to max_batch_size (the commented line) skips that padding, - # which pays under concurrency (requests at different steps, one or two outside the dit). - # H200, 1024^2, 8 steps, 2 rounds (before the fused q/k norm + rotary): + # The capture budget per size is len(capture_caption_lengths) x len(capture_batch_sizes) graphs. + # A step of 3 or 5-7 rows would replay the next captured size with dummy rows; at 16 concurrent + # requests steps average 5.7 rows, so with only 1/2/4/8 captured about 30% of a step's rows would + # be dummies. Capturing every batch size pays under concurrency (requests at different steps, one + # or two outside the dit). H200, 1024^2, 8 steps, 2 rounds (before the fused q/k norm + rotary): # c=4 1.05-1.08 -> 1.22 images/s, c=8 1.12-1.14 -> 1.21-1.22, # c=16 1.17-1.19 -> 1.23, c=1 unchanged, same peak VRAM. But each graph is a compile of its own: - # 32 graphs start ~100 s slower than 16 on a warm Inductor cache, ~8 min slower on a cold one. + # 32 graphs start ~100 s slower than 16 on a warm Inductor cache, ~8 min slower on a cold one, + # and the 64 graphs of the two sizes cost 5.3 GiB of VRAM and take a warm start from about 96 s to about 220 s on MI325X. # The text encoder captures its own geometric_batch_sizes: powers of two below max_batch_size, # then max_batch_size (4 graphs here, any caption length). capture_caption_lengths: [32, 64, 96, 128] - capture_batch_sizes: [1, 2, 4, 8] - # capture_batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8] # for deployments that batch 3, 5, 6 or 7 + # every batch size up to max_batch_size; [1, 2, 4, 8] halves the graphs for deployments that start cold often + capture_batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8] max_batch_size: 8 # compiled VAE decode (max-autotune), not bit-exact with the eager decode: before the folded upsample and split # GroupNorm it was 3x faster at 1024^2 and ~56 dB from eager, and the autotuner may pick other conv kernels in another diff --git a/docs/architecture.rst b/docs/architecture.rst index 190497ecbe..2107ba0df3 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -209,6 +209,14 @@ step. ``os.environ.setdefault`` in the conductor before it spawns workers, so values you export win. ROCm turns off MIOpen's split-K conv solver there (see :doc:`environment_variables`); the other platforms set nothing. +- **Compile pool.** ``Platform.compile_pool_overrides()`` gives + ``torch._inductor.config`` values in force while a worker builds Inductor's compile + pool, before it touches the model, since a submodule can compile at load. The pool reads + them once, so they stay out of every graph's cache key, which the same setting in the + environment would not (that recompiles every graph a golden compile cache holds). ROCm + turns off the pool's idle quiesce: the sidecar + retires a pool without waiting for it, and a worker of the next pool forked meanwhile can + hang forever on the old pool's shutdown lock, along with the compile that gave it a job. - **Capture refusal.** A resource whose kernel can't run some capture bucket inside a graph returns a reason from ``Resource.capture_refusal(spec)``. Both graph runners ask before capturing, and a refused bucket runs eagerly. Answer from the spec and the diff --git a/mstar/model/z_image/submodules.py b/mstar/model/z_image/submodules.py index 77183dac88..4ea3211336 100644 --- a/mstar/model/z_image/submodules.py +++ b/mstar/model/z_image/submodules.py @@ -23,11 +23,9 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo +from mstar.engine import recompile_limit from mstar.engine.accelerator_graph_config import geometric_batch_sizes -from mstar.model.components.diffusion.compile_utils import ( - compile_transformer_forward, - step_schedule_recompile_limit, -) +from mstar.model.components.diffusion.compile_utils import compile_transformer_forward from mstar.model.components.diffusion.decode_utils import ( VAE_MAX_DECODE_BATCH, compile_vae_decode, @@ -145,9 +143,16 @@ def __init__( if module is not None: compile_transformer_forward( module, eager_rounding=compile_eager_rounding, exact_ops=compile_exact_ops, - recompile_limit=step_schedule_recompile_limit(step_dtypes), + recompile_limit=self._recompile_limit(step_dtypes), ) + def _recompile_limit(self, step_dtypes) -> int: + # every block runs one code object, so dynamo's per-code limit counts its shapes over all three spans of every + # captured graph, at most two per graph; a step schedule compiles each dtype on its own + per_dtype = max(recompile_limit(), 2 * len(self.capture_buckets) * len(self.capture_batch_sizes)) + dtypes = 1 if step_dtypes is None else len(step_dtypes.dtypes) + return min(per_dtype * dtypes, 256 * dtypes) + def bucket_key_for(self, fwd_info: CurrentForwardPassInfo) -> ZShape: return shape_from_metadata(self.config, fwd_info.step_metadata) diff --git a/mstar/platforms/compile.py b/mstar/platforms/compile.py index 022af68b5f..a2de694090 100644 --- a/mstar/platforms/compile.py +++ b/mstar/platforms/compile.py @@ -1,4 +1,4 @@ -"""The platform's inductor settings, applied to one compiled callable.""" +"""The platform's inductor settings, applied to one compiled callable or to the compile pool.""" from __future__ import annotations @@ -39,6 +39,23 @@ def forward(*args, **kwargs): return forward +def start_compile_pool() -> None: + """Build Inductor's compile pool under the platform's ``compile_pool_overrides``. + + Call before the process's first compile, since the first caller builds the + one pool. The pool reads these when built, and keeps them after the patch. + A value set in the environment would stay in the config instead, and so in + every graph's cache key. + """ + overrides = current_platform.compile_pool_overrides() + if not overrides: + return + + import torch._inductor.config as inductor_config + from torch._inductor.async_compile import AsyncCompile + + with inductor_config.patch(**overrides): + AsyncCompile.warm_pool() @functools.cache def compiled(fn: Callable[..., Any], dynamic: bool | None = None) -> Callable[..., Any]: """``torch.compile(fn)`` under the platform's inductor overrides, built on diff --git a/mstar/platforms/interface.py b/mstar/platforms/interface.py index a5239519ec..f1ce5b5369 100644 --- a/mstar/platforms/interface.py +++ b/mstar/platforms/interface.py @@ -41,6 +41,18 @@ class HardwareQueueColocation(WorkerColocation): def start(self) -> None: logger.info("mps: true on ROCm: nothing to start, the GPU schedules the workers' queues itself") +BF16_HEADER_FLAG = "-include hip/hip_bf16.h" + + +def _with_bf16_header(flags: str | None) -> str: + """``FLASHINFER_EXTRA_CUDAFLAGS`` with the header amd-flashinfer 0.6.18+amd.1's + AITER prefill sources need: their generated config names ``__hip_bfloat16`` + before anything declares it, since FlashInfer's ``-DUSE_ROCM=1`` turns off + the include in torch's ``BFloat16.h``. Later releases include it themselves.""" + if flags is None: + return BF16_HEADER_FLAG + return flags if BF16_HEADER_FLAG in flags else f"{flags} {BF16_HEADER_FLAG}" + class Platform: """One accelerator stack; ``current_platform`` is the one this process runs. @@ -113,6 +125,12 @@ def inductor_overrides( """ return {} + def compile_pool_overrides(self) -> dict[str, object]: + """``torch._inductor.config`` values in force while a worker builds + Inductor's compile pool. Unlike ``inductor_overrides`` these never reach + a graph's cache key. Empty keeps inductor's defaults.""" + return {} + def process_env_defaults(self) -> dict[str, str | Callable[[str | None], str]]: """Environment variables for settings a runtime or library reads once at init. Workers apply them with ``os.environ.setdefault`` before their @@ -305,8 +323,16 @@ def inductor_overrides( overrides["post_grad_custom_post_pass"] = post return overrides - def process_env_defaults(self) -> dict[str, str]: + def compile_pool_overrides(self) -> dict[str, object]: + # Inductor's sidecar retires a pool that returned nothing for a minute + # without waiting for it. A worker of the next pool forked meanwhile + # inherits the old pool's shutdown lock held by a thread it doesn't + # have, and hangs at its first full GC with the first job it took. + return {"quiesce_async_compile_pool": False} + + def process_env_defaults(self) -> dict[str, str | Callable[[str | None], str]]: return { + "FLASHINFER_EXTRA_CUDAFLAGS": _with_bf16_header, # MIOpen's split-K implicit-GEMM conv solvers are not deterministic. Turning # them off costs about 1% and makes same-seed outputs repeat. The backward # one runs transposed convs (LTX-2.5's vocoder upsamplers). diff --git a/mstar/worker/engine_manager.py b/mstar/worker/engine_manager.py index 62568c5516..ce4be11382 100644 --- a/mstar/worker/engine_manager.py +++ b/mstar/worker/engine_manager.py @@ -16,6 +16,7 @@ from mstar.graph.runtime.base import GraphRuntime from mstar.model.base import Model, node_moe_parallel_configs, parse_autocast_dtype, serving_resources from mstar.model.submodule_base import NodeSubmodule +from mstar.platforms.compile import start_compile_pool from mstar.utils.device import DeviceCapability from mstar.utils.streams import reset_device_scheduling @@ -269,6 +270,10 @@ def build( the latter, so there is no separate cache config. The conductor's ``validate_deployment`` has already checked the specs and overrides. """ + # Before the model does anything: the process's first compile builds the + # pool, and a submodule can compile at load (klein warms its VAE decode). + if (device_capability or DeviceCapability.of(device)).is_available: + start_compile_pool() specs = serving_resources(model, model_config) # YAML autocast_dtype over the model's own, but never re-enabling diff --git a/test/modular/test_platforms.py b/test/modular/test_platforms.py index a1d51f4f9c..b62beaaef1 100644 --- a/test/modular/test_platforms.py +++ b/test/modular/test_platforms.py @@ -1,9 +1,13 @@ """``mstar.platforms``: which hardware the process runs on.""" import copy +import importlib.util import multiprocessing import os import pickle +import shutil +import subprocess +import sys from types import SimpleNamespace import pytest @@ -11,6 +15,7 @@ from mstar import platforms from mstar.platforms import CpuPlatform, CudaPlatform, RocmPlatform, XpuPlatform, current_platform +from mstar.platforms.interface import BF16_HEADER_FLAG @pytest.fixture @@ -163,6 +168,60 @@ def test_a_callable_env_default_adds_to_the_value_already_set(monkeypatch): assert os.environ["PLATFORM_ENV_TEST_SET"] == "user-default" +@pytest.mark.parametrize( + ("exported", "expected"), + [ + (None, BF16_HEADER_FLAG), + ("-DUSER_FLAG", f"-DUSER_FLAG {BF16_HEADER_FLAG}"), + (f"{BF16_HEADER_FLAG} -DUSER_FLAG", f"{BF16_HEADER_FLAG} -DUSER_FLAG"), + ], +) +def test_rocm_adds_the_bf16_header_to_flashinfer_flags(monkeypatch, exported, expected): + monkeypatch.setattr(torch.version, "hip", "7.0") + monkeypatch.setenv("FLASHINFER_EXTRA_CUDAFLAGS", "") + if exported is None: + monkeypatch.delenv("FLASHINFER_EXTRA_CUDAFLAGS") + else: + monkeypatch.setenv("FLASHINFER_EXTRA_CUDAFLAGS", exported) + + # the conductor applies the defaults, then each worker again + platforms.apply_process_env_defaults() + platforms.apply_process_env_defaults() + assert os.environ["FLASHINFER_EXTRA_CUDAFLAGS"] == expected + + +# aiter's import asks rocminfo for a live GPU; the module generator reads only its path +_BF16_PREFILL_JIT = """ +import importlib.util, sys, types +import torch +aiter = types.ModuleType("aiter") +aiter.__file__ = importlib.util.find_spec("aiter").origin +sys.modules["aiter"] = aiter +from flashinfer.jit.rocm.modules import gen_batch_prefill_module +bf16 = torch.bfloat16 +gen_batch_prefill_module("aiter", bf16, bf16, bf16, torch.int32, 64, 64, 0, False, False, False).build() +""" + + +def test_flashinfer_builds_its_bf16_aiter_prefill_with_the_rocm_flags(tmp_path): + if importlib.util.find_spec("aiter") is None: + pytest.skip("needs AITER") + pytest.importorskip("flashinfer.jit.rocm.modules") + if shutil.which("amdclang++") is None: + pytest.skip("needs amdclang++") + env = { + **os.environ, + "FLASHINFER_EXTRA_CUDAFLAGS": BF16_HEADER_FLAG, + "FLASHINFER_ROCM_ARCH_LIST": "gfx942", + "FLASHINFER_WORKSPACE_BASE": str(tmp_path), + } + built = subprocess.run( + [sys.executable, "-c", _BF16_PREFILL_JIT], + env=env, capture_output=True, text=True, check=False, + ) + assert built.returncode == 0, built.stderr[-2000:] + built.stdout[-2000:] + + def test_a_worker_sets_its_env_defaults_before_anything_else(monkeypatch): from mstar.conductor import conductor @@ -246,6 +305,74 @@ def test_only_rocm_turns_off_miopens_split_k_conv_solvers(solver): assert solver not in platform.process_env_defaults() +def test_only_rocm_builds_inductors_compile_pool_without_its_quiesce(): + assert RocmPlatform().compile_pool_overrides() == {"quiesce_async_compile_pool": False} + for platform in (CudaPlatform(), XpuPlatform(), CpuPlatform()): + assert platform.compile_pool_overrides() == {} + + +def test_the_compile_pool_keeps_its_settings_out_of_the_cache_key(): + # a real sidecar, in a process of its own since torch keeps the first pool it builds + code = """ +import torch._inductor.config as config +from torch._inductor.async_compile import AsyncCompile +from mstar.platforms import RocmPlatform, compile as platform_compile + +key = config.save_config_portable() +platform_compile.current_platform = RocmPlatform() +platform_compile.start_compile_pool() +pool = AsyncCompile.process_pool() +assert pool.timer is None, "the pool would still quiesce" +assert config.save_config_portable() == key, "a graph's cache key moved" +pool.shutdown() +""" + subprocess.run([sys.executable, "-c", code], env={**os.environ, "TORCHINDUCTOR_COMPILE_THREADS": "2"}, check=True) + + +@pytest.mark.parametrize( + ("platform", "available", "built"), + [("RocmPlatform", True, True), ("RocmPlatform", False, False), ("CudaPlatform", True, False)], +) +def test_a_worker_builds_the_compile_pool_before_it_touches_the_model(platform, available, built): + # a model can compile at load, and the first compile builds the pool under inductor's defaults + code = f""" +from types import SimpleNamespace + +import torch +from torch._inductor.async_compile import AsyncCompile + +from mstar import platforms +from mstar.platforms import compile as platform_compile +from mstar.worker.engine_manager import EngineManager + + +class ModelTouched(Exception): + pass + + +class Model: + def __getattr__(self, name): + raise ModelTouched + + +platform_compile.current_platform = platforms.{platform}() +try: + EngineManager.build( + node_names={{"node"}}, device=torch.device("cpu"), model_config={{}}, parallel_groups=None, + transfer_engine_info=None, model=Model(), graph_runtime=None, + device_capability=SimpleNamespace(is_available={available}), + ) +except ModelTouched: + pass +assert AsyncCompile.process_pool.cache_info().currsize == {int(built)} +if {built}: + pool = AsyncCompile.process_pool() + assert pool.timer is None, "the pool would still quiesce" + pool.shutdown() +""" + subprocess.run([sys.executable, "-c", code], env={**os.environ, "TORCHINDUCTOR_COMPILE_THREADS": "2"}, check=True) + + def test_platform_knob_defaults_are_empty_but_for_rocm(): assert RocmPlatform().knob_defaults() == {"dit_exact_ops": "0"} for platform in (CudaPlatform(), XpuPlatform(), CpuPlatform()): diff --git a/test/modular/test_z_image_model.py b/test/modular/test_z_image_model.py index 188c813e75..ad33678b88 100644 --- a/test/modular/test_z_image_model.py +++ b/test/modular/test_z_image_model.py @@ -116,6 +116,31 @@ def test_capture_shapes_cover_sizes_times_caption_buckets(): assert _make_model(accelerator_graph=False).capture_buckets() == [] +@pytest.mark.parametrize(("capture_sizes", "batch_sizes", "limit"), [ + ([[1024, 1024]], [1, 2, 4, 8], None), # 2 x 4 x 4 graphs fit the process-wide limit + ([[1024, 1024], [768, 1344]], list(range(1, 9)), 128), +]) +def test_the_transformers_recompile_limit_covers_the_capture_plan(monkeypatch, capture_sizes, batch_sizes, limit): + from mstar.engine import recompile_limit + from mstar.model.components.diffusion import compile_utils + from mstar.model.z_image.submodules import ZImageDenoiseSubmodule + + seen = {} + + def fake_compile(fn, **kwargs): + seen.update(kwargs) + return fn + + monkeypatch.setattr(compile_utils.torch, "compile", fake_compile) + model = _make_model(capture_sizes=capture_sizes, capture_caption_lengths=[32, 64, 96, 128], + capture_batch_sizes=batch_sizes) + ZImageDenoiseSubmodule( + torch.nn.Linear(2, 2), ZImageConfig(), loop_name=DENOISE_LOOP, attn_resource_key=None, + capture_buckets=model.capture_buckets(), capture_batch_sizes=model.capture_batch_sizes, + ) + assert seen["recompile_limit"] == (limit or recompile_limit()) + + def test_postprocess_png_and_adapter(): model = _make_model() assert model.postprocess(torch.zeros(1, 3, 16, 16, dtype=torch.uint8), "image")[:4] == b"\x89PNG" diff --git a/test/rocm/test_varlen_parity_shapes.py b/test/rocm/test_varlen_parity_shapes.py index 7d50f04147..fdb47186bb 100644 --- a/test/rocm/test_varlen_parity_shapes.py +++ b/test/rocm/test_varlen_parity_shapes.py @@ -346,10 +346,11 @@ def test_patch_embed_conv(case, worker_env): cin, cout, kernel, shape, dtype, tol = _PATCH_CONVS[case] env = dict(os.environ) - for key in _built_for().process_env_defaults(): + defaults = _built_for().process_env_defaults() + for key in defaults: env.pop(key, None) if worker_env: - env.update(_built_for().process_env_defaults()) + env.update({k: v(None) if callable(v) else v for k, v in defaults.items()}) res = subprocess.run( [sys.executable, "-c", _CONV_CHECK, json.dumps([cin, cout, kernel, shape, dtype])], env=env, capture_output=True, text=True, timeout=300, check=True,