diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 5ec4323..e68a2df 100644 --- a/magi_compiler/_api.py +++ b/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 diff --git a/magi_compiler/cuda/cudart.py b/magi_compiler/cuda/cudart.py index aac558a..b8017ae 100644 --- a/magi_compiler/cuda/cudart.py +++ b/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}") diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 48843b5..3e7f0f9 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/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,7 +286,7 @@ 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) @@ -293,26 +294,27 @@ def fix_graph_device_placement(module: torch.nn.Module): 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) diff --git a/magi_compiler/offload/offload_warpper.py b/magi_compiler/offload/offload_warpper.py index 4f2c788..02453ef 100644 --- a/magi_compiler/offload/offload_warpper.py +++ b/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 diff --git a/magi_compiler/offload/profiler.py b/magi_compiler/offload/profiler.py index 4a01478..1ed07b5 100644 --- a/magi_compiler/offload/profiler.py +++ b/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,10 +49,10 @@ 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() @@ -56,7 +60,7 @@ def get_h2d_bandwidth(self, size_mb=1024, iters=3, warmup=3, dtype=torch.float32 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: diff --git a/magi_compiler/offload/scheduler.py b/magi_compiler/offload/scheduler.py index c3569b4..6706cd1 100644 --- a/magi_compiler/offload/scheduler.py +++ b/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,7 +455,7 @@ 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 @@ -460,8 +463,8 @@ def prefetch(self, current_node_name: str, ctx: OffloadRuntimeContext): 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) diff --git a/magi_compiler/utils/device.py b/magi_compiler/utils/device.py index ebcd246..23fda5c 100644 --- a/magi_compiler/utils/device.py +++ b/magi_compiler/utils/device.py @@ -42,3 +42,29 @@ def device_capability(device: int = 0) -> Tuple[int, int]: def device_capability_major(device: int = 0) -> int: """Convenience wrapper: just the major-capability int (0 if no CUDA).""" return device_capability(device)[0] + + +def accelerator_module(): + """``torch.cuda`` on NVIDIA GPUs, the out-of-tree backend module (e.g. ``torch_musa``) elsewhere.""" + import torch as _torch + + return _torch.get_device_module() + + +def accelerator_type() -> str: + """Device type of the current accelerator, e.g. ``"cuda"`` or ``"musa"``.""" + import torch as _torch + + acc = _torch.accelerator.current_accelerator() + return acc.type if acc is not None else "cuda" + + +def is_host_scalar(val) -> bool: + """0-dim CPU tensor, e.g. a Python float lifted by Dynamo. + + Device kernels accept these as-is, and they must stay on the host: a fake one moved to the device loses + its constant value, so a later ``.item()`` becomes data-dependent. + """ + import torch as _torch + + return isinstance(val, _torch.Tensor) and val.dim() == 0 and val.device.type == "cpu" diff --git a/magi_compiler/utils/nvtx.py b/magi_compiler/utils/nvtx.py index f5729fb..2d4cd5d 100644 --- a/magi_compiler/utils/nvtx.py +++ b/magi_compiler/utils/nvtx.py @@ -24,12 +24,14 @@ # global var for torch.autograd.profiler.emit_nvtx _EMIT_NVTX_CTX: None | torch.autograd.profiler.emit_nvtx = None +_HAS_NVTX = hasattr(torch._C, "_nvtx") @torch.library.custom_op("magi_compile::nvtx_range_push", mutates_args=()) def nvtx_range_push(event_name: str) -> None: """torch.ops.magi_compile.nvtx_range_push""" - torch.cuda.nvtx.range_push(event_name) + if _HAS_NVTX: + torch.cuda.nvtx.range_push(event_name) @nvtx_range_push.register_fake @@ -40,7 +42,8 @@ def _(event_name: str) -> None: @torch.library.custom_op("magi_compile::nvtx_range_pop", mutates_args=()) def nvtx_range_pop() -> None: """torch.ops.magi_compile.nvtx_range_pop""" - torch.cuda.nvtx.range_pop() + if _HAS_NVTX: + torch.cuda.nvtx.range_pop() @nvtx_range_pop.register_fake @@ -67,14 +70,14 @@ def __enter__(self): # nor modifying a variable not in the current scope # so we have no choice but assign a constant event name when compiling nvtx_range_push("torch compile region") - else: + elif _HAS_NVTX: torch.cuda.nvtx.range_push(self.enter_name) return self def __exit__(self, *excinfo): if torch.compiler.is_compiling(): nvtx_range_pop() - else: + elif _HAS_NVTX: torch.cuda.nvtx.range_pop() diff --git a/tests/feature_tests/test_cpu_offload.py b/tests/feature_tests/test_cpu_offload.py index 3064b11..23a3c2d 100644 --- a/tests/feature_tests/test_cpu_offload.py +++ b/tests/feature_tests/test_cpu_offload.py @@ -264,5 +264,41 @@ def forward(self, x): # assert model_mtd.mlp.up_proj.weight.device.type == "cpu" +@magi_register_custom_op("athena::scale_by", infer_output_meta_fn=["x"]) +def scale_by(x: torch.Tensor, scale: float) -> torch.Tensor: + return x * scale + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA support") +def test_cpu_offload_float_attribute(device, mlp_config): + """ + A float module attribute is lifted by Dynamo into a 0-dim CPU tensor input and read back with .item() + for the custom op; the offload path must keep it on the host. + """ + with set_cpu_offload(True): + + @magi_compile(dynamic_arg_dims={"x": 0}) + class ScaledMLP(torch.nn.Module): + def __init__(self, config: MLPConfig): + super().__init__() + self.scale = 0.5 + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False, dtype=config.params_dtype) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False, dtype=config.params_dtype) + + def forward(self, x): + return self.down_proj(scale_by(F.silu(self.up_proj(x)), self.scale)) + + model = TransformerWrapper(mlp_config, mlp_cls=ScaledMLP) + model.cuda() + assert model.mlp.up_proj.weight.device.type == "cpu" + + with torch.no_grad(): + for num_tokens in (4, 8, 16): + x = torch.randn(num_tokens, mlp_config.hidden_size, device=device, dtype=mlp_config.params_dtype) + up, down = model.mlp.up_proj.weight.to(device), model.mlp.down_proj.weight.to(device) + expected = model.attention_proj(my_attention(*[F.linear(F.silu(F.linear(x, up)) * 0.5, down)] * 3)) + torch.testing.assert_close(model(x), expected, rtol=2e-2, atol=2e-2) + + if __name__ == "__main__": pytest.main([__file__, "-v"])