fix: _patch_cpu_offload_apply handles .to(cuda) in addition to .cuda() - #59
Open
cennn wants to merge 1 commit into
Open
fix: _patch_cpu_offload_apply handles .to(cuda) in addition to .cuda()#59cennn wants to merge 1 commit into
cennn wants to merge 1 commit 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.
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)