[WIP]: feat(offload): run model_cpu_offload on non-CUDA accelerators (MUSA)

magi_compiler/_api.py

@@ -31,6 +31,7 @@
 from magi_compiler.magi_backend.magi_compiler_base import MagiCompileState
 from magi_compiler.utils import compilation_counter, envs, magi_logger
 from magi_compiler.utils.compile_time_monitor import CompileMonitor
+from magi_compiler.utils.device import accelerator_type
 from magi_compiler.utils.host_memory import fmt_host_mem
 
 from .config import CompileConfig, CompileMode
@@ -694,7 +695,7 @@ def _cpu_apply(self, fn):
         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
+                is_to_cuda = fn(torch.empty(0, device="cpu")).device.type == accelerator_type()
             except Exception:
                 pass
 

magi_compiler/cuda/cudart.py

@@ -18,7 +18,10 @@
 
 import torch
 
+from magi_compiler.utils.device import accelerator_type
+
 _cudart = None
+_musart = None
 
 
 def init_cudart():
@@ -40,23 +43,39 @@ def init_cudart():
     return None
 
 
+def init_musart():
+    global _musart
+    if _musart is None:
+        try:
+            _musart = ctypes.CDLL("libmusart.so")
+        except OSError:
+            pass
+    return _musart
+
+
 def pin_memory_in_place(tensor: torch.Tensor):
     """
-    Pin memory in-place using cudaHostRegister.
+    Pin memory in-place using cudaHostRegister (musaHostRegister on MUSA).
     """
-    if tensor.is_cuda:
+    if tensor.device.type != "cpu":
         return tensor
-    cudart = init_cudart()
-    if cudart is None:
+    # Pick by accelerator, not by which library loads: a MUSA host may also ship libcudart.
+    if accelerator_type() == "musa":
+        lib = init_musart()
+        host_register = lib.musaHostRegister if lib is not None else None
+    else:
+        lib = init_cudart()
+        host_register = lib.cudaHostRegister if lib is not None else None
+    if host_register is None:
         return tensor
 
     ptr = tensor.data_ptr()
     size = tensor.numel() * tensor.element_size()
-    res = cudart.cudaHostRegister(ctypes.c_void_p(ptr), ctypes.c_size_t(size), 0)
+    res = host_register(ctypes.c_void_p(ptr), ctypes.c_size_t(size), 0)
 
     # 0: success, 712: already registered
     # https://www.autodesk.com/support/technical/article/caas/sfdcarticles/sfdcarticles/GPU-and-CUDA-call-failed-error-messages-with-Arnold-5-3-5-and-Maya-2024.html
     if res == 0 or res == 712:
         return tensor
     else:
-        raise RuntimeError(f"cudaHostRegister failed with error code {res}")
+        raise RuntimeError(f"{host_register.__name__} failed with error code {res}")

magi_compiler/magi_backend/magi_backend.py

@@ -35,6 +35,7 @@
 from magi_compiler.offload.offload_warpper import OffloadWrapper
 from magi_compiler.passes import CustomJointGraphPartitionFn, FullGraphPassManager, PostGradPassManager, pass_context
 from magi_compiler.utils import compilation_counter, compute_code_hash, compute_graph_hash, compute_hash, magi_logger
+from magi_compiler.utils.device import accelerator_module, accelerator_type, is_host_scalar
 from magi_compiler.utils.visualize import save_fx_graph_visualization
 
 from ._cache_data_cls import CacheEntry, CacheHandle
@@ -271,12 +272,12 @@ def _device_is_cpu(val) -> bool:
     return isinstance(val, str) and val == 'cpu'
 
 
-def _recursive_to_device(val, target_device: int):
+def _recursive_to_device(val, target_device: torch.device):
     """Recursively move tensor-like *val* (or nested list/tuple) to *target_device*."""
     if isinstance(val, (list, tuple)):
         items = [_recursive_to_device(v, target_device) for v in val]
         return type(val)(items) if any(n is not o for n, o in zip(items, val)) else val
-    if hasattr(val, 'device') and str(val.device) == 'cpu':
+    if hasattr(val, 'device') and str(val.device) == 'cpu' and not is_host_scalar(val):
         new_val = val.to(target_device)
         if isinstance(val, torch.nn.Parameter):
             new_val = torch.nn.Parameter(new_val, requires_grad=val.requires_grad)
@@ -285,34 +286,35 @@ def _recursive_to_device(val, target_device: int):
 
 
 def fix_graph_device_placement(module: torch.nn.Module):
-    """Rewrite CPU device refs and example_values to CUDA in an FX graph."""
+    """Rewrite CPU device refs and example_values to the accelerator in an FX graph."""
     for _, child in module.named_children():
         fix_graph_device_placement(child)
 
     if not isinstance(module, torch.fx.GraphModule):
         return
 
     needs_recompile = False
-    target_device = torch.cuda.current_device()
+    target_device = torch.device(accelerator_type(), accelerator_module().current_device())
 
     for node in module.graph.nodes:
         if node.op == 'call_function' and 'device' in node.kwargs:
             if _device_is_cpu(node.kwargs['device']):
-                node.update_kwarg('device', torch.device('cuda', target_device))
+                node.update_kwarg('device', target_device)
                 needs_recompile = True
 
-        if node.op == 'call_method' and node.target == 'to':
+        # torch_musa wraps Tensor.to, so on MUSA Dynamo records the wrapped method as orig_to.
+        if node.op == 'call_method' and node.target in ('to', 'orig_to'):
             new_args = list(node.args)
             changed = False
             for i, arg in enumerate(new_args):
                 if _device_is_cpu(arg):
-                    new_args[i] = torch.device('cuda', target_device)
+                    new_args[i] = target_device
                     changed = True
             if changed:
                 node.args = tuple(new_args)
                 needs_recompile = True
             if 'device' in node.kwargs and _device_is_cpu(node.kwargs['device']):
-                node.update_kwarg('device', torch.device('cuda', target_device))
+                node.update_kwarg('device', target_device)
                 needs_recompile = True
 
     cpu_fix_count = 0
@@ -327,7 +329,7 @@ def fix_graph_device_placement(module: torch.nn.Module):
             cpu_fix_count += 1
 
     if needs_recompile:
-        magi_logger.info('[fix_device] fixed %d CPU example_values to cuda:%s', cpu_fix_count, target_device)
+        magi_logger.info('[fix_device] fixed %d CPU example_values to %s', cpu_fix_count, target_device)
         module.recompile()
 
 
@@ -364,8 +366,8 @@ def run(self, *args):
         if self.compile_config.offload_config.model_cpu_offload:
             fix_graph_device_placement(self.module)
             for i, arg in enumerate(fake_args):
-                if isinstance(arg, torch.Tensor):
-                    fake_args[i] = arg.cuda()
+                if isinstance(arg, torch.Tensor) and not is_host_scalar(arg):
+                    fake_args[i] = arg.to(accelerator_type())
 
         with self.fake_mode, enable_python_dispatcher():
             return super().run(*fake_args)

magi_compiler/offload/offload_warpper.py

@@ -23,6 +23,7 @@
 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.device import accelerator_module, accelerator_type, is_host_scalar
 from magi_compiler.utils.nvtx import add_nvtx_event
 
 from ..magi_depyf.timeline import observe_lifecycle
@@ -33,8 +34,10 @@ 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.device_module = accelerator_module()
+        self.device_type = accelerator_type()
+        self.compute_stream = self.device_module.current_stream()
+        self.h2d_stream = self.device_module.Stream()
 
         self.warmup = True
         self.second_call = False
@@ -105,15 +108,15 @@ def _prepare_inputs(self, args) -> Dict[Node, Any]:
 
             # case 1: input tensor
             if not is_weight:
-                if isinstance(arg_val, torch.Tensor):
-                    arg_val = arg_val.to("cuda", non_blocking=False)
+                if isinstance(arg_val, torch.Tensor) and not is_host_scalar(arg_val):
+                    arg_val = arg_val.to(self.device_type, 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
+                    t = arg_val.to(self.device_type, 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
@@ -122,7 +125,7 @@ def _prepare_inputs(self, args) -> Dict[Node, Any]:
             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)
+                    arg_val = arg_val.to(self.device_type, non_blocking=False)
                 elif not self.warmup:
                     if node in self.submod_0_weight_handoff:
                         arg_val = self.submod_0_weight_handoff[node]
@@ -166,7 +169,7 @@ def __call__(self, *args):
                     self.profiler.start_compute_profile(node.name, self.compute_stream)
 
                 with add_nvtx_event(node.name):
-                    with torch.cuda.stream(self.compute_stream):
+                    with self.device_module.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)
@@ -183,7 +186,7 @@ def __call__(self, *args):
                     parent_node, idx = node.args
                     env[node] = env[parent_node][idx]
                 else:
-                    with torch.cuda.stream(self.compute_stream):
+                    with self.device_module.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)
@@ -204,7 +207,7 @@ def __call__(self, *args):
                 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:
+                        if isinstance(tensor_obj, torch.Tensor) and tensor_obj.device.type != "cpu":
                             tensor_obj.record_stream(self.compute_stream)
                         del env[input_node]
         return None

magi_compiler/offload/profiler.py

@@ -16,26 +16,30 @@
 
 import torch
 
+from magi_compiler.utils.device import accelerator_module, accelerator_type
+
 
 class OffloadProfiler:
     def __init__(self):
+        self.device_module = accelerator_module()
         self.compute_events: Dict[str, Dict[str, torch.cuda.Event]] = {}
         self.timings: Dict[str, Dict[str, float]] = {}
 
     def start_compute_profile(self, name: str, stream: torch.cuda.Stream):
         if name not in self.compute_events:
             self.compute_events[name] = {}
-        start_event = torch.cuda.Event(enable_timing=True)
+        start_event = self.device_module.Event(enable_timing=True)
         start_event.record(stream)
         self.compute_events[name]["start"] = start_event
 
     def end_compute_profile(self, name: str, stream: torch.cuda.Stream):
-        end_event = torch.cuda.Event(enable_timing=True)
+        end_event = self.device_module.Event(enable_timing=True)
         end_event.record(stream)
         self.compute_events[name]["end"] = end_event
 
-    def get_h2d_bandwidth(self, size_mb=1024, iters=3, warmup=3, dtype=torch.float32, device=torch.device("cuda")):
-        torch.cuda.synchronize()
+    def get_h2d_bandwidth(self, size_mb=1024, iters=3, warmup=3, dtype=torch.float32, device=None):
+        device = device or accelerator_type()
+        self.device_module.synchronize()
 
         num_elements = size_mb * 1024 * 1024 // torch.tensor([], dtype=dtype).element_size()
 
@@ -45,18 +49,18 @@ def get_h2d_bandwidth(self, size_mb=1024, iters=3, warmup=3, dtype=torch.float32
         # warmup
         for _ in range(warmup):
             gpu_tensor.copy_(cpu_tensor, non_blocking=True)
-        torch.cuda.synchronize()
+        self.device_module.synchronize()
 
-        start_event = torch.cuda.Event(enable_timing=True)
-        end_event = torch.cuda.Event(enable_timing=True)
+        start_event = self.device_module.Event(enable_timing=True)
+        end_event = self.device_module.Event(enable_timing=True)
 
         start_event.record()
 
         for _ in range(iters):
             gpu_tensor.copy_(cpu_tensor, non_blocking=True)
 
         end_event.record()
-        torch.cuda.synchronize()
+        self.device_module.synchronize()
 
         elapsed_ms = start_event.elapsed_time(end_event)
         elapsed_s = elapsed_ms / 1000.0
@@ -72,7 +76,7 @@ def broadcast_obj(self, obj, src=0):
         return obj_list[0]
 
     def summarize(self) -> Dict[str, Dict[str, float]]:
-        torch.cuda.synchronize()
+        self.device_module.synchronize()
         results = {}
         for name, evs in self.compute_events.items():
             if name not in results:

magi_compiler/offload/scheduler.py

@@ -22,6 +22,7 @@
 
 from magi_compiler.config import CompileConfig
 from magi_compiler.utils import magi_logger
+from magi_compiler.utils.device import accelerator_module, accelerator_type
 
 
 @dataclass
@@ -86,6 +87,8 @@ def __init__(
         self.weight_sizes = weight_sizes
 
         self.kept_weights: Set[str] = set()
+        self.device_module = accelerator_module()
+        self.device_type = accelerator_type()
 
     @abc.abstractmethod
     def schedule_kept_weights(self, profile_data: Dict[str, float]):
@@ -189,7 +192,7 @@ def prefetch(self, current_node_name: str, ctx: OffloadRuntimeContext):
 
         weight_names = self.submod_weights_map.get(target_node.name, [])
 
-        with torch.cuda.stream(h2d_stream):
+        with self.device_module.stream(h2d_stream):
             for w_name in weight_names:
                 if w_name in self.kept_weights:
                     continue
@@ -200,10 +203,10 @@ def prefetch(self, current_node_name: str, ctx: OffloadRuntimeContext):
 
                 if is_next_iter:
                     if w_name in buffers:
-                        submod_0_next_iter_weights[w_node] = buffers[w_name].to("cuda", non_blocking=True)
+                        submod_0_next_iter_weights[w_node] = buffers[w_name].to(self.device_type, non_blocking=True)
 
                 elif w_node in env and hasattr(env[w_node], 'device') and env[w_node].device.type == "cpu":
-                    env[w_node] = env[w_node].to("cuda", non_blocking=True)
+                    env[w_node] = env[w_node].to(self.device_type, non_blocking=True)
 
 
 @SchedulerFactory.register("COST_EFFECTIVE")
@@ -270,9 +273,9 @@ def __init__(
         # prefetch schedule: { trigger_node_name : target_submod_index }
         # when execute the trigger_node_name, prefetch the target_submod_index
         self.prefetch_schedule = collections.defaultdict()
-        self.submod_load_events: Dict[str, torch.cuda.Event] = collections.defaultdict(torch.cuda.Event)
+        self.submod_load_events: Dict[str, torch.cuda.Event] = collections.defaultdict(self.device_module.Event)
         for i in range(len(self.submod_nodes)):
-            self.submod_load_events[self.submod_nodes[i].name] = torch.cuda.Event()
+            self.submod_load_events[self.submod_nodes[i].name] = self.device_module.Event()
 
     def schedule_kept_weights(self, profile_data: Dict[str, Any]):
         """
@@ -443,7 +446,7 @@ def prefetch(self, current_node_name: str, ctx: OffloadRuntimeContext):
         if h2d_idx is None:
             return
 
-        compute_complete_event = torch.cuda.Event()
+        compute_complete_event = self.device_module.Event()
         compute_complete_event.record(compute_stream)
         h2d_stream.wait_event(compute_complete_event)
 
@@ -452,16 +455,16 @@ def prefetch(self, current_node_name: str, ctx: OffloadRuntimeContext):
 
         weight_names = self.submod_weights_map.get(h2d_node.name, [])
 
-        with torch.cuda.stream(h2d_stream):
+        with self.device_module.stream(h2d_stream):
             for w_name in weight_names:
                 if w_name in self.kept_weights:
                     continue
 
                 w_node = self.name_node_map.get(w_name)
 
                 if w_node in env and hasattr(env[w_node], 'device') and env[w_node].device.type == "cpu":
-                    env[w_node] = env[w_node].to("cuda", non_blocking=True)
+                    env[w_node] = env[w_node].to(self.device_type, non_blocking=True)
                 elif is_last_step and w_name in buffers:
-                    submod_0_next_iter_weights[w_node] = buffers[w_name].to("cuda", non_blocking=True)
+                    submod_0_next_iter_weights[w_node] = buffers[w_name].to(self.device_type, non_blocking=True)
 
             self.submod_load_events[h2d_node.name].record(h2d_stream)