fix: keep accelerate offload hook when swapping Linear -> SCLinear - #7
Merged
Allenjin123 merged 1 commit intoJun 2, 2026
Merged
Conversation
With device_map="auto" + CPU offload, the original nn.Linear carries an accelerate AlignDevicesHook whose weights_map holds the real CPU weights while the module param is a 'meta' placeholder. replace_linears_with_sc dropped that hook, so the SCLinear's weight stayed meta and the SC forward crashed with 'RuntimeError: Tensor on device meta is not on the expected device'. Transfer the existing hook to the replacement module so offloaded weights are materialized at forward. Verified: Llama-3.1-70B SC runs on a single GPU via CPU offload, bit-identical output to the 2-GPU sharded run. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes replace_linears_with_sc() so it preserves Hugging Face Accelerate’s device dispatch / CPU offload behavior when swapping nn.Linear modules to SCLinear, preventing “meta tensor” crashes when loading large models with device_map="auto" and CPU offload.
Changes:
- Detects an existing Accelerate
_hf_hookon the originalnn.Linearand transfers it to the newSCLinear. - Keeps the prior
.to(device, dtype)behavior only for non-hooked (non-dispatched/offloaded) modules.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Allenjin123
approved these changes
Jun 2, 2026
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
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.
Problem
With
device_map="auto"+ CPU offload, the originalnn.Linearcarries an accelerateAlignDevicesHookwhoseweights_mapholds the real CPU weights while the module's own param is ametaplaceholder.replace_linears_with_scdropped that hook when swapping inSCLinear, so the weight stayedmetaand the SC forward crashed:This blocked running 70B SC on a single GPU (where the model must offload to fit).
Fix
Transfer the existing
_hf_hookfrom the original module to the newSCLinear(remove from the discarded module first), so offloadedmetaweights are materialized to the execution device at forward. Non-offloaded modules (no hook) are unchanged.Validation
Llama-3.1-70B SC runs on a single GPU via CPU offload, producing bit-identical output to the 2-GPU sharded run.
Pairs with the multi-GPU device-guard fix in scmp_kernels (CrucibleComputingGroup/scmp_kernels#20).
🤖 Generated with Claude Code