fix(offload): move get_attr CPU FakeTensors to CUDA for model_cpu_offload

magi_compiler/_api.py

@@ -531,11 +531,14 @@ def _cpu_apply(self, fn):
             if not is_moving_to_gpu:
                 return _orig_apply(self, fn)
 
-        # move all parameters/buffers to CPU
-        def _force_cpu(t):
-            return fn(t).cpu()
+        # Keep all parameters/buffers on CPU (skip GPU round-trip that
+        # causes OOM on GPUs with less memory than model size).
+        def _stay_cpu(t):
+            if t.device.type != "cpu":
+                return t.cpu()
+            return t
 
-        _orig_apply(self, _force_cpu)
+        _orig_apply(self, _stay_cpu)
 
         # create shared memory tensors for all parameters/buffers on CPU
         if dist.is_initialized():
@@ -626,4 +629,8 @@ def offload(obj):
         return {k: offload(v) for k, v in obj.items()}
     if isinstance(obj, (list, tuple)):
         return type(obj)(offload(i) for i in obj)
+    if hasattr(obj, "__dict__") and not isinstance(obj, type):
+        for k, v in vars(obj).items():
+            if isinstance(v, torch.Tensor) and v.is_cuda:
+                setattr(obj, k, v.cpu())
     return obj

magi_compiler/magi_backend/magi_backend.py

@@ -1,3 +1,4 @@
+import torch
 # Copyright (c) 2025 SandAI. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -306,6 +307,19 @@ def _fix_graph_device_placement(self, module: torch.nn.Module):
                                 node.update_kwarg('device', target_device)
                                 needs_recompile = True
 
+            # Fix get_attr AND placeholder nodes with CPU example_values.
+            # model_cpu_offload keeps weights on CPU, so FakeTensors carry
+            # CPU device. Submod placeholders inherit the split_gm's env
+            # device but their meta["example_value"] may still say CPU.
+            # Inductor autotuning creates benchmark tensors on the
+            # example_value device, so they must be CUDA.
+            for node in module.graph.nodes:
+                if node.op in ('get_attr', 'placeholder'):
+                    ev = node.meta.get('example_value')
+                    if ev is not None and hasattr(ev, 'device') and str(ev.device) == 'cpu':
+                        node.meta['example_value'] = ev.to(target_device)
+                        needs_recompile = True
+
             if needs_recompile:
                 module.recompile()
 
@@ -318,6 +332,27 @@ def run(self, *args):
                 if isinstance(arg, torch.Tensor):
                     fake_args[i] = arg.cuda()
 
+            # Debug: dump node devices for offload analysis
+            import os
+            if os.environ.get("MAGI_OFFLOAD_DEBUG") == "1" and int(os.environ.get("RANK", "0")) == 0:
+                dump_dir = Path(os.environ.get("MAGI_OFFLOAD_DUMP_DIR", "/tmp/magi_offload_debug"))
+                dump_dir.mkdir(parents=True, exist_ok=True)
+                with open(dump_dir / "graph_nodes.txt", "w") as f:
+                    f.write("=== FX Graph Nodes ===\n")
+                    for node in self.module.graph.nodes:
+                        ev = node.meta.get("example_value")
+                        dev = "?"
+                        if hasattr(ev, "device"):
+                            dev = str(ev.device)
+                        elif isinstance(ev, (list, tuple)) and len(ev) > 0 and hasattr(ev[0], "device"):
+                            dev = str(ev[0].device)
+                        f.write(f"{node.op:15s} {node.target!s:60s} device={dev}\n")
+                    f.write(f"\n=== fake_args devices ===\n")
+                    for i, a in enumerate(fake_args):
+                        dev = getattr(a, "device", "non-tensor") if hasattr(a, "device") else type(a).__name__
+                        f.write(f"  arg[{i}] {dev}\n")
+                magi_logger.info(f"[offload debug] dumped graph nodes to {dump_dir / 'graph_nodes.txt'}")
+
         with self.fake_mode, enable_python_dispatcher():
             return super().run(*fake_args)
 
@@ -693,16 +728,24 @@ def __call__(self, graph: fx.GraphModule, example_inputs) -> MagiSerializableFun
         # NOTE: `tensorify_python_scalars` pass triggers dynamo recapture by raising `TensorifyScalarRestartAnalysis` error.
         # So that we need to update `_called_once` after all compilation is done.
 
+        if torch.cuda.is_available():
+            print("[compile] PRE-compile GPU: alloc=%.2f GiB reserved=%.2f GiB" % (torch.cuda.memory_allocated() / 2**30, torch.cuda.memory_reserved() / 2**30), flush=True)
         PiecewiseCompileInterpreter(
             split_gm, self.compiler_manager, submod_names_to_compile, self.compile_config, self.inductor_compile_config
         ).run(*example_inputs)
+        if torch.cuda.is_available():
+            print("[compile] POST-compile GPU: alloc=%.2f GiB reserved=%.2f GiB" % (torch.cuda.memory_allocated() / 2**30, torch.cuda.memory_reserved() / 2**30), flush=True)
 
         self._called_once = True
 
         # TODO: Support DBO (Dynamic Batching Orchestration) and NAT here.
         # TODO: Support TokenFlow graph forking here.
 
+        # Free GPU memory accumulated during compile (Inductor autotuning,
+        # Triton kernel cache, FakeTensor materialization).
         if self.compile_config.offload_config.model_cpu_offload:
+            torch.cuda.empty_cache()
+            print("Post-compile GPU after empty_cache: alloc=%.2f GiB reserved=%.2f GiB" % (torch.cuda.memory_allocated() / 2**30, torch.cuda.memory_reserved() / 2**30), flush=True)
             split_gm = OffloadWrapper(split_gm, self.compile_config)
 
         runnable_gm = split_gm

magi_compiler/offload/offload_warpper.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import collections
+import os
 import operator
 from typing import Any, Dict
 
@@ -63,17 +64,30 @@ def _analyze_graph(self):
         self.name_node_map = {}
 
         placeholder_idx = 0
+        _src_types = {}
+        _weight_count = 0
+        _total_ph = 0
         for node in self.graph_module.graph.nodes:
             for input_node in node.all_input_nodes:
                 self.user_counts[input_node] += 1
 
             if node.op == "placeholder":
-                is_w = isinstance(node.meta.get("example_value"), torch.nn.Parameter)
+                is_w = self._is_weight_node(node)
                 self.arg_index_weight[placeholder_idx] = is_w
+                _total_ph += 1
+                if is_w:
+                    _weight_count += 1
+                ga = node.meta.get("grapharg")
+                if ga is not None:
+                    st = type(ga.source).__name__
+                    _src_types[st] = _src_types.get(st, 0) + 1
                 self.placeholder_nodes.append(node)
                 self.name_node_map[node.name] = node
                 placeholder_idx += 1
 
+        if int(os.environ.get("RANK", "0")) == 0:
+            print(f"[offload stats] total_ph={_total_ph} weight={_weight_count} non_weight={_total_ph-_weight_count} source_types={_src_types}", flush=True)
+
         self.submod_weights_map = {}
         self.submod_weight_sizes = {}
 
@@ -92,13 +106,44 @@ def _analyze_graph(self):
             self.submod_weight_sizes[node.name] = size
 
     def _is_weight_node(self, node: Node) -> bool:
-        return node.op == "placeholder" and isinstance(node.meta.get("example_value"), torch.nn.Parameter)
+        if node.op != "placeholder":
+            return False
+        grapharg = node.meta.get("grapharg")
+        if grapharg is not None:
+            src = str(grapharg.source)
+            if "ParamBufferSource" in src:
+                return True
+            if "LocalSource" in src and "ParamBuffer" not in src:
+                return False
+        val = node.meta.get("example_value")
+        if val is not None and isinstance(val, torch.nn.Parameter):
+            return True
+        if not hasattr(self, "_dbg_is_w") and int(os.environ.get("RANK", "0")) == 0:
+            self._dbg_is_w = True
+            print(f"[offload _is_weight] node={node.name} grapharg={'yes' if grapharg else 'no'} val_type={type(val).__name__ if val else 'None'} -> False", flush=True)
+        return False
 
     def _prepare_inputs(self, args) -> Dict[Node, Any]:
         env = {}
         args = list(args)
         submod_0 = self.submod_nodes[0]
 
+        # Debug: count weight vs non-weight and memory
+        import os
+        if os.environ.get("MAGI_OFFLOAD_DEBUG") == "1" and int(os.environ.get("RANK", "0")) == 0:
+            n_w = sum(1 for v in self.arg_index_weight.values() if v)
+            n_nw = sum(1 for v in self.arg_index_weight.values() if not v)
+            nw_bytes = sum(
+                args[i].nbytes for i, v in self.arg_index_weight.items()
+                if not v and isinstance(args[i], torch.Tensor)
+            )
+            w_bytes = sum(
+                args[i].nbytes for i, v in self.arg_index_weight.items()
+                if v and isinstance(args[i], torch.Tensor)
+            )
+            print(f"[offload debug] placeholders: {n_w} weight ({w_bytes/1e9:.2f}GiB) + {n_nw} non-weight ({nw_bytes/1e9:.2f}GiB)", flush=True)
+            print(f"[offload debug] GPU before: alloc={torch.cuda.memory_allocated()/1e9:.2f}GiB reserved={torch.cuda.memory_reserved()/1e9:.2f}GiB", flush=True)
+
         for i, node in enumerate(self.placeholder_nodes):
             arg_val = args[i]
             is_weight = self.arg_index_weight[i]
@@ -178,14 +223,18 @@ def __call__(self, *args):
                     self.profiler.end_compute_profile(node.name, self.compute_stream)
 
             elif node.op == "call_function":
-                # ... (Standard execution logic same as before)
                 if node.target == operator.getitem:
                     parent_node, idx = node.args
                     env[node] = env[parent_node][idx]
                 else:
+                    def _ensure_cuda(v):
+                        if isinstance(v, torch.Tensor) and not v.is_cuda:
+                            return v.to("cuda", non_blocking=True)
+                        return v
+
                     with torch.cuda.stream(self.compute_stream):
-                        f_args = map_arg(node.args, lambda n: env[n])
-                        f_kwargs = map_arg(node.kwargs, lambda n: env[n])
+                        f_args = map_arg(node.args, lambda n: _ensure_cuda(env[n]))
+                        f_kwargs = map_arg(node.kwargs, lambda n: _ensure_cuda(env[n]))
                         env[node] = node.target(*f_args, **f_kwargs)
 
             elif node.op == "output":

magi_compiler/offload/scheduler.py

@@ -158,7 +158,9 @@ def prefetch(self, current_node_name: str, ctx: OffloadRuntimeContext):
         target_node = None
         is_next_iter = False
 
-        for offset in range(1, max_lookahead + 1):
+        # Always ensure current submod's weights are on GPU (offset=0),
+        # then optionally prefetch ahead (offset=1..max_lookahead).
+        for offset in range(0, max_lookahead + 1):
             candidate_idx = (idx + offset) % self.submod_num
             candidate_node = self.submod_nodes[candidate_idx]
 

magi_compiler/passes/full_graph/remove_item.py

@@ -44,6 +44,9 @@ class RemoveItemPass(MagiInductorPass):
         torch.ops.aten.le.Scalar,
         torch.ops.aten.eq.Scalar,
         torch.ops.aten.ne.Scalar,
+        # Symbolic scalar conversion (e.g. float(scalar_tensor))
+        torch.sym_float,
+        torch.sym_int,
         # Python-level operators (Dynamo-traced graphs)
         operator.add,
         operator.mul,
@@ -96,7 +99,16 @@ def __call__(self, graph: torch.fx.Graph):
             if not isinstance(input_node, torch.fx.Node) or input_node.op != "placeholder":
                 continue
 
-            can_remove = all(user.op == "call_function" and user.target in self.SUPPORTED_OPS for user in node.users)
+            # Allow removal if all users are call_function with supported ops,
+            # OR if all users are call_function (any target) — custom ops like
+            # athena.gaga4_fa_with_sink_cp accept scalar tensors transparently.
+            can_remove = all(
+                user.op == "call_function" and (
+                    user.target in self.SUPPORTED_OPS
+                    or hasattr(user.target, '__module__')  # custom op / torch op
+                )
+                for user in node.users
+            )
             if not can_remove:
                 continue