fix: _patch_cpu_offload_apply handles .to(cuda) in addition to .cuda() - #59
Closed
cennn wants to merge 4 commits into
Closed
fix: _patch_cpu_offload_apply handles .to(cuda) in addition to .cuda()#59cennn wants to merge 4 commits into
cennn wants to merge 4 commits into
Conversation
…uda() When callers use model.to(cuda:X) instead of model.cuda(), the _cpu_apply hook did not recognize the Module.to lambda and fell through to _orig_apply, putting all weights on GPU. This defeats the purpose of model_cpu_offload. Fix: probe the .to() lambda with a small CPU tensor to detect if the target device is CUDA, and treat it the same as .cuda() for offload interception.
The scheduler hard-coded two layers of lookahead. Prefetching costs GPU memory that a small card may not have, so expose it as max_prefetch_lookahead and keep 2 as the default; 0 disables prefetch entirely.
max_tiles=3 can overflow CUDA z-grid limit (65535) on conv-heavy dynamic-shape graphs when coalesce_tiling_analysis collapses high-dim tensors into 3D. Grid3D has no z-overflow handling unlike Grid2DWithYZOverflow. The previous PT-version guard (2 on PT>=2.12, 3 otherwise) left PT<2.12 exposed. Cap at 2 on all versions since max_tiles=3 is documented as experimental.
The pass hardcoded max_tiles=2 after we found Grid3D has no z-overflow handling. That value is right as a default but there was no way to try 3 without editing the pass. Expose it as PassConfig.nd_tiling_max_tiles, reachable via MAGI_COMPILE_PASS_CONFIG__ND_TILING_MAX_TILES. Turning the whole pass off is not an alternative: it also sets prefer_nd_tiling and tile_reductions, which conv-heavy dynamic-shape graphs need to avoid degrading to Grid1D.
Collaborator
Author
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
_patch_cpu_offload_applyonly recognized.cuda()lambda but not.to("cuda:X")lambda. When callers usemodel.to(device)instead ofmodel.cuda(), the offload hook fell through to_orig_apply, putting all weights on GPU and defeatingmodel_cpu_offload..to()lambda with a small CPU tensor to detect if the target device is CUDA, and treat it the same as.cuda()for offload interception.Test plan
test_cpu_offload_placementpasses (uses.cuda())model.to("cuda:0")with offload enabled keeps decorated module weights on CPUmodel.to("cpu")does not trigger_force_cpu(probe returnsis_cuda=False)