From 68205660111304e8d43b10563b1e0b113a77d301 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Tue, 18 Aug 2026 11:48:45 +0800 Subject: [PATCH 1/2] fix: _patch_cpu_offload_apply now handles .to(cuda) in addition to .cuda() When callers use model.to(cuda:X) instead of model.cuda(), the _cpu_apply hook did not recognize the Module.to lambda and fell through to _orig_apply, putting all weights on GPU. This defeats the purpose of model_cpu_offload. Fix: probe the .to() lambda with a small CPU tensor to detect if the target device is CUDA, and treat it the same as .cuda() for offload interception. --- magi_compiler/_api.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 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 From d8cbfd12476e88f97339097288fde55d63fbccfb Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 19 Aug 2026 10:27:31 +0800 Subject: [PATCH 2/2] feat(offload): make the prefetch lookahead configurable The scheduler hard-coded two layers of lookahead. Prefetching costs GPU memory that a small card may not have, so expose it as max_prefetch_lookahead and keep 2 as the default; 0 disables prefetch entirely. --- magi_compiler/config.py | 1 + magi_compiler/offload/scheduler.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/magi_compiler/config.py b/magi_compiler/config.py index f0206ac..d5592d8 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -180,6 +180,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