MagiCompiler/magi_compiler/offload/offload_warpper.py

# Copyright (c) 2026 SandAI. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import collections
import operator
from typing import Any, Dict

import torch
from torch.fx import GraphModule, Node
from torch.fx.node import map_arg

from magi_compiler.config import CompileConfig
from magi_compiler.offload.profiler import OffloadProfiler
from magi_compiler.offload.scheduler import OffloadRuntimeContext, SchedulerFactory
from magi_compiler.utils.nvtx import add_nvtx_event

from ..magi_depyf.timeline import observe_lifecycle


class OffloadExecutor:
    def __init__(self, graph_module: GraphModule, compile_config: CompileConfig):
        self.graph_module = graph_module
        self.compile_config = compile_config

        self.compute_stream = torch.cuda.current_stream()
        self.h2d_stream = torch.cuda.Stream()

        self.warmup = True
        self.second_call = False
        self.buffers: Dict[str, torch.Tensor] = {}
        self.persistent_weights: Dict[str, torch.Tensor] = {}
        self.submod_0_weight_handoff: Dict[Node, torch.Tensor] = {}

        self._analyze_graph()
        self.profiler = OffloadProfiler()

        common_args = {
            "submod_nodes": self.submod_nodes,
            "submod_weights_map": self.submod_weights_map,
            "name_node_map": self.name_node_map,
            "weight_sizes": self.submod_weight_sizes,
        }

        self.scheduler = SchedulerFactory.create(self.compile_config, common_args)

    def _analyze_graph(self):
        self.submod_nodes = [n for n in self.graph_module.graph.nodes if n.op == "call_module"]

        self.placeholder_nodes = []
        self.arg_index_weight = {}
        self.user_counts = collections.defaultdict(int)
        self.name_node_map = {}

        placeholder_idx = 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)
                self.arg_index_weight[placeholder_idx] = is_w
                self.placeholder_nodes.append(node)
                self.name_node_map[node.name] = node
                placeholder_idx += 1

        self.submod_weights_map = {}
        self.submod_weight_sizes = {}

        for node in self.submod_nodes:
            weight_names = []
            size = 0
            for arg in node.args:
                if isinstance(arg, Node) and self._is_weight_node(arg):
                    if arg.name in self.name_node_map:
                        weight_names.append(arg.name)
                        val = arg.meta.get("example_value")
                        if val is not None:
                            size += val.numel() * val.element_size()

            self.submod_weights_map[node.name] = weight_names
            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)

    def _prepare_inputs(self, args) -> Dict[Node, Any]:
        env = {}
        args = list(args)
        submod_0 = self.submod_nodes[0]

        for i, node in enumerate(self.placeholder_nodes):
            arg_val = args[i]
            is_weight = self.arg_index_weight[i]

            # case 1: input tensor
            if not is_weight:
                if isinstance(arg_val, torch.Tensor):
                    arg_val = arg_val.to("cuda", non_blocking=False)
                env[node] = arg_val
                continue

            # case 2: kept weight
            if self.scheduler.is_kept(node.name):
                if node.name not in self.persistent_weights:
                    t = arg_val.to("cuda", non_blocking=False) if arg_val.device.type == "cpu" else arg_val
                    self.persistent_weights[node.name] = t
                env[node] = self.persistent_weights[node.name]
                continue

            # case 3: submod 0 weight
            if submod_0 in node.users:
                if self.warmup and arg_val.device.type == "cpu":
                    self.buffers[node.name] = arg_val
                    arg_val = arg_val.to("cuda", non_blocking=False)
                elif not self.warmup:
                    if node in self.submod_0_weight_handoff:
                        arg_val = self.submod_0_weight_handoff[node]
                        del self.submod_0_weight_handoff[node]

            env[node] = arg_val

        return env

    def _finalize_warmup(self):
        profile_results = self.profiler.summarize()
        self.scheduler.schedule_kept_weights(profile_results)
        self.warmup = False

    def __call__(self, *args):
        if len(self.submod_nodes) == 0:
            return self.graph_module(*args)

        env = self._prepare_inputs(args)
        current_user_counts = self.user_counts.copy()
        runtime_ctx = OffloadRuntimeContext(
            env=env,
            h2d_stream=self.h2d_stream,
            compute_stream=self.compute_stream,
            buffers=self.buffers,
            submod_0_handoff=self.submod_0_weight_handoff,
            need_profile=self.second_call or self.warmup,
        )
        need_profile = self.second_call

        for node in self.graph_module.graph.nodes:
            if node.op == "placeholder":
                continue

            elif node.op == "call_module":
                self.scheduler.prefetch(node.name, runtime_ctx)

                if need_profile:
                    if torch.distributed.is_initialized():
                        torch.distributed.barrier()
                    self.profiler.start_compute_profile(node.name, self.compute_stream)

                with add_nvtx_event(node.name):
                    with torch.cuda.stream(self.compute_stream):
                        s_args = map_arg(node.args, lambda n: env[n])
                        s_kwargs = map_arg(node.kwargs, lambda n: env[n])
                        env[node] = getattr(self.graph_module, node.target)(*s_args, **s_kwargs)
                        del s_args, s_kwargs

                if need_profile:
                    if torch.distributed.is_initialized():
                        torch.distributed.barrier()
                    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:
                    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])
                        env[node] = node.target(*f_args, **f_kwargs)

            elif node.op == "output":
                if self.second_call:
                    self._finalize_warmup()
                    self.second_call = False
                if self.warmup:
                    self.second_call = True
                    self.warmup = False

                return map_arg(node.args[0], lambda n: env[n])

            # Memory Management
            for input_node in node.all_input_nodes:
                current_user_counts[input_node] -= 1
                if current_user_counts[input_node] == 0:
                    if input_node in env:
                        tensor_obj = env[input_node]
                        if isinstance(tensor_obj, torch.Tensor) and tensor_obj.is_cuda:
                            tensor_obj.record_stream(self.compute_stream)
                        del env[input_node]
        return None


class OffloadWrapper:
    @observe_lifecycle("offload_wrap")
    def __init__(self, graph_module: torch.fx.GraphModule, compile_config: CompileConfig):
        self.executor = OffloadExecutor(graph_module, compile_config)

    def __call__(self, *args):
        return self.executor(*args)