Repository navigation
Conversation
The offload path hard-coded torch.cuda, "cuda", .is_cuda and libcudart, so it failed on MUSA unless torchada remapped torch.cuda. Use the current accelerator's device module instead; on CUDA this is torch.cuda, so behaviour there is unchanged. - utils/device.py: accelerator_module() / accelerator_type(). - offload executor, schedulers, profiler: streams, events, synchronize and H2D copies go through them. - _api.py: detect .to(<accelerator>) by device type rather than .is_cuda. - magi_backend.py: fix_graph_device_placement and the offload fake args target the accelerator. - cuda/cudart.py: pin_memory_in_place falls back to libmusart musaHostRegister when libcudart is absent. - utils/nvtx.py: skip NVTX ranges when torch is built without NVTX. 8x S5000 (torch 2.9.0, no torchada), toy 4-layer block, ratio 0/0.5, BASE/COST_EFFECTIVE/HEURISTIC: weights stay on CPU and pinned, rel err vs eager 4.2e-3 (bf16); 2 ranks: giant buffers pinned, 1.24 ms/call vs 4.23 ms unpinned. H200: test_cpu_offload, test_ep_shared_memory, test_fix_to_cpu_in_graph, test_shm_memory_peak, cache/test_cache_topology_isolation and api_tests pass (140).
# Conflicts: # magi_compiler/magi_backend/magi_backend.py
Dynamo lifts unspecialized Python floats (e.g. a module's eps) into 0-dim CPU tensor inputs and reads them back with .item() where a Python float is required (custom op float args, Triton launches). The offload path moved every tensor input to the accelerator: the fake copy loses its constant, so the piecewise fake run raised DataDependentOutputException, and at runtime the scalar was a device tensor. Leave 0-dim CPU tensors alone in the offload fake args, fix_graph_device_placement and the executor's input preparation; device kernels accept CPU scalars as before. test_cpu_offload_float_attribute fails without the fix (DataDependentOutputException) and passes with it on H200; the H200 offload/cache/api suites still pass (147). Found on MAGI-2 400B, where every RMSNorm eps and mHC scale hit this once the Triton paths were traced.
torch_musa replaces Tensor.to with a wrapper that calls self.orig_to, so on MUSA Dynamo records call_method orig_to. fix_graph_device_placement only matched 'to', so an in-graph .to(device=x.device) traced with the offload's CPU inputs stayed device_put(..., cpu), and a Triton kernel then received a CPU pointer at runtime. MAGI-2 400B W8A8 on 8x S5000, gpu_resident_weight_ratio=0.0: compiles and runs, 5.15 s/step vs 5.08 s/step without offload, peak device memory 21.7 GB/card vs 73.8 GB.
pin_memory_in_place tried libcudart first and used cudaHostRegister whenever it loaded, so a MUSA host that also ships libcudart called the CUDA runtime and raised. Select musaHostRegister or cudaHostRegister from accelerator_type() instead. Also rename the offload executor/scheduler/profiler attributes rt -> device_module and device -> device_type (it holds the type string, not a torch.device).
cennn
changed the title
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
🗂️ PR Category
📝 Description
model_cpu_offloadhard-coded CUDA (torch.cudastreams/events,.to("cuda"),.is_cuda,cudaHostRegister,torch.cuda.nvtx), so it could not run on other PyTorch accelerator backends such as Moore Threads MUSA. Offload matters most there: large models that do not fit a card's memory are exactly what we want to run on these devices.1. Offload runs on the current accelerator, CUDA unchanged.
Streams, events,
synchronize, H2D copies and the offload hook's.to(<device>)detection go throughtorch.get_device_module()/torch.accelerator.current_accelerator()(on CUDA these aretorch.cuda/cuda, so CUDA behaviour is identical). Pinning picksmusaHostRegisterorcudaHostRegisterby accelerator type. Two backend quirks are handled: NVTX ranges are skipped when torch is built without NVTX, andfix_graph_device_placementalso rewritesorig_to(torch_musa wrapsTensor.to, so Dynamo recordscall_method orig_to; without this an in-graph.to(device=x.device)stayed on CPU and a Triton kernel received a CPU pointer).2. Fix: float module attributes broke offload compilation (CUDA too).
Dynamo lifts unspecialized Python floats (a module's
eps,scale, ...) into 0-dim CPU tensor inputs and reads them back with.item()where a Python float is required (custom-op float args, Triton launches). Offload moved every tensor input to the device, so the fake copy lost its constant and the piecewise fake run raisedDataDependentOutputException. 0-dim CPU tensors now stay on the host (offload fake args, graph device placement, executor input preparation). Newtest_cpu_offload_float_attributefails before and passes after.3. Real-model result on 8× S5000. A 400B MoE video DiT (W8A8), 540p 10 s i2v, 4 steps,
gpu_resident_weight_ratio=0.0:Outputs match the non-offload run by eye:
Small model on S5000 (1 and 2 ranks,
BASE/COST_EFFECTIVE/HEURISTIC, ratio 0 / 0.5): weights stay on CPU and are pinned, relative error vs eager 4.2e-3 (bf16); with 2 ranks the pinned per-rank buffer takes 1.24 ms/call vs 4.23 ms unpinned.Other
x.device.type == "cuda"/x.is_cudabranch in the traced model takes its CPU path, silently and for good. On the model above this dropped its Triton fused kernels and cost ~80% per step until the model accounted for it.is_available()):symm_mem/,cuda_graph_mgr.py,profiling/runtime_estimator.py, EVT fusion.