fix(offload): move get_attr CPU FakeTensors to CUDA for model_cpu_offload - #62
Closed
cennn wants to merge 5 commits into
Closed
fix(offload): move get_attr CPU FakeTensors to CUDA for model_cpu_offload#62cennn wants to merge 5 commits into
cennn wants to merge 5 commits into
Conversation
Three fixes for running gaga4 base model with model_cpu_offload: 1. _fix_graph_device_placement: move get_attr CPU FakeTensors to CUDA (weights stay CPU but compile needs CUDA FakeTensors) 2. offload(): recurse into object __dict__ to CPU-offload tensor attributes of non-tensor args (e.g. ModalityDispatcher.permute_mapping) 3. RemoveItemPass: accept any call_function user for placeholder-sourced .item() nodes, not just the SUPPORTED_OPS whitelist. This handles custom ops like athena.gaga4_fa_with_sink_cp that receive scalar args.
…evice + RemoveItem
_fix_graph_device_placement moves placeholder example_values to CUDA via .to(device), which drops nn.Parameter type from FakeTensor. This caused _is_weight_node to classify all 1303 weights as inputs, leading to _prepare_inputs moving 46 GiB to GPU and OOM. Fix: check grapharg.source for ParamBufferSource (Dynamo provenance) instead of relying on isinstance(val, nn.Parameter).
call_function nodes (custom ops like gaga4_mh_moe) can reference weight/buffer placeholders. The scheduler only prefetches weights for call_module nodes, so these remain on CPU causing device mismatch. Fix: _ensure_cuda in call_function arg mapping.
Collaborator
Author
|
Superseded by #64 which includes a broader device-placement fix (all node types, not just get_attr). |
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.
Summary
_fix_graph_device_placementonly fixed factory functions (torch.empty/zerosetc.) but missedget_attrFX nodes. Whenmodel_cpu_offload=True,_patch_cpu_offload_applykeeps model weights on CPU, so Dynamo-inlinedget_attrnodes carry CPUexample_value. DuringPiecewiseCompileInterpreter.run(), inputs are moved to CUDA butget_attrweights stay CPU →index_select(CUDA, CPU)device conflict in Inductor.Fix: extend
_fix_graph_device_placementto also moveget_attrnodemeta["example_value"]from CPU to CUDA.Context
Found while trying to run Gaga4 base model on 8x5090 (SM120) with
model_cpu_offload. The model has extensiveindex_selectwith modality dispatching, which hit this device mismatch.Test plan
model_cpu_offload=Trueon 5090