diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 7bf5f21..0f65384 100644 --- a/magi_compiler/_api.py +++ b/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 diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 239b224..e74fa41 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/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 diff --git a/magi_compiler/offload/offload_warpper.py b/magi_compiler/offload/offload_warpper.py index 4f2c788..6ec9e6c 100644 --- a/magi_compiler/offload/offload_warpper.py +++ b/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": diff --git a/magi_compiler/offload/scheduler.py b/magi_compiler/offload/scheduler.py index c3569b4..ebc897f 100644 --- a/magi_compiler/offload/scheduler.py +++ b/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] diff --git a/magi_compiler/passes/full_graph/remove_item.py b/magi_compiler/passes/full_graph/remove_item.py index e86b743..8ee47ff 100644 --- a/magi_compiler/passes/full_graph/remove_item.py +++ b/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