From 0a5881cc97c414c6a59ed3e854b3c9f485e36cc2 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Fri, 21 Aug 2026 20:19:01 +0800 Subject: [PATCH] fix: offload .to(cuda) compat + configurable nd_tiling max_tiles + prefetch lookahead --- magi_compiler/_api.py | 15 +++++++++++++-- magi_compiler/config.py | 13 +++++++++++++ magi_compiler/offload/scheduler.py | 2 +- .../piecewise_graph/nd_tiling_workaround.py | 11 ++++++----- 4 files changed, 33 insertions(+), 8 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index e7c6209..7bf5f21 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -509,6 +509,17 @@ def _cpu_apply(self, fn): id_cpu_lambda = getattr(fn, "__qualname__", "") == "Module.cpu.." is_to_lambda = getattr(fn, "__qualname__", "") == "Module.to..convert" + # Detect .to("cuda:X") by probing the lambda on a small CPU tensor. + # .to("cpu") also produces is_to_lambda=True but should not trigger offload. + is_to_cuda = False + if is_to_lambda and not getattr(self, "_magi_offloaded_once", False): + try: + is_to_cuda = fn(torch.empty(0, device="cpu")).is_cuda + except Exception: + pass + + is_moving_to_gpu = is_cuda_lambda or is_to_cuda + # after first time to call _apply(cuda), skip "Module.to" and "Module.cpu" and "Module.cuda" if getattr(self, "_magi_offloaded_once", False): if is_cuda_lambda or id_cpu_lambda or is_to_lambda: @@ -516,8 +527,8 @@ def _cpu_apply(self, fn): else: return _orig_apply(self, fn) else: - # first time to call _apply(cuda), move all parameters/buffers to CPU - if not is_cuda_lambda: + # first time to call _apply(cuda) or _apply(to_cuda), move all parameters/buffers to CPU + if not is_moving_to_gpu: return _orig_apply(self, fn) # move all parameters/buffers to CPU diff --git a/magi_compiler/config.py b/magi_compiler/config.py index f0206ac..181c94a 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -92,6 +92,18 @@ class PassConfig(BaseModel): "Env var: MAGI_COMPILE_PASS_CONFIG__ENABLE_ND_TILING_WORKAROUND (1/0/true/false)." ), ) + nd_tiling_max_tiles: int = Field( + 2, + ge=1, + le=3, + description=( + "max_tiles the ND-tiling workaround sets. 2 (default) is safe: Inductor's Grid2D " + "folds a y-grid overflow into z. 3 is experimental -- Grid3D has no z-overflow " + "handling, so a conv-heavy dynamic-shape graph (turbo VAE at 1080p) can exceed " + "CUDA's 65535 z-grid limit. " + "Env var: MAGI_COMPILE_PASS_CONFIG__ND_TILING_MAX_TILES." + ), + ) enable_mm_epilogue_fusion: bool = Field( False, description=( @@ -180,6 +192,7 @@ class OffloadConfig(BaseModel): OffloadPolicy.COST_EFFECTIVE, description="The policy for offloading the model to CPU." ) bandwidth_safety_factor: float = Field(0.9, description="The safety factor for the H2D bandwidth.") + max_prefetch_lookahead: int = Field(2, description="Max layers to prefetch ahead. 0 disables prefetch to save GPU memory.") class FSDPConfig(BaseModel): diff --git a/magi_compiler/offload/scheduler.py b/magi_compiler/offload/scheduler.py index 566a4ea..c3569b4 100644 --- a/magi_compiler/offload/scheduler.py +++ b/magi_compiler/offload/scheduler.py @@ -154,7 +154,7 @@ def prefetch(self, current_node_name: str, ctx: OffloadRuntimeContext): except ValueError: return - max_lookahead = 2 + max_lookahead = self.compile_config.offload_config.max_prefetch_lookahead target_node = None is_next_iter = False diff --git a/magi_compiler/passes/piecewise_graph/nd_tiling_workaround.py b/magi_compiler/passes/piecewise_graph/nd_tiling_workaround.py index a69c32d..21de6d1 100644 --- a/magi_compiler/passes/piecewise_graph/nd_tiling_workaround.py +++ b/magi_compiler/passes/piecewise_graph/nd_tiling_workaround.py @@ -18,7 +18,6 @@ import torch from ...magi_depyf.timeline import emit_pass_lifecycle -from ...utils.envs import IS_PT_212 from ..pass_base import MagiInductorPass @@ -40,7 +39,9 @@ def __call__(self, graph: torch.fx.Graph): torch._inductor.config.triton.prefer_nd_tiling = True torch._inductor.config.triton.tile_reductions = True - # PT 2.12 Inductor generates invalid 3D-grid reduction kernels with - # max_tiles=3 (program_id(2) mapped to a non-existent grid dim). - # Cap at 2 on PT >= 2.12 until the upstream fix lands. - torch._inductor.config.triton.max_tiles = 2 if IS_PT_212 else 3 + # Capped at 2 by default: Grid3D has no z-overflow handling, so a + # conv-heavy dynamic-shape graph can exceed CUDA's 65535 z-grid limit. + # Raise it via MAGI_COMPILE_PASS_CONFIG__ND_TILING_MAX_TILES to test 3. + from ...config import get_compile_config + + torch._inductor.config.triton.max_tiles = get_compile_config().pass_config.nd_tiling_max_tiles