[Offload] Support views - #786
Open
kylesayrs wants to merge 6 commits into
Open
Conversation
Contributor
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughChangesOffload view-aware caching
Pretrained loading initialization
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
coderabbitai
Bot
added
the
documentation
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/compressed_tensors/offload/cache/base.py (2)
124-129: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winUse
__setitem__to properly populate view metadata and reference counts.
from_mappingcurrently assigns directly toinstance.offloaded_values, bypassing__setitem__. This meansview_indexandref_counterare never populated during module initialization. As a result, subsequent onloading of views will fail (by erroneously returning the full base tensor) and disk cache cleanup will break due to missing reference counts.Refactoring this to use
instance[name] = tensorensures all tracking logic executes correctly. As per path instructions, verify that offload hooks interact correctly with model loading.♻️ Proposed fix
- instance.offloaded_values = { - name: instance.offload(tensor) for name, tensor in mapping.items() - } + for name, tensor in mapping.items(): + instance[name] = tensor🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/compressed_tensors/offload/cache/base.py` around lines 124 - 129, Update from_mapping to populate the instance through __setitem__ by assigning each mapping entry with instance[name] = tensor instead of directly replacing offloaded_values. Preserve the existing offload behavior while ensuring __setitem__ records view_index and ref_counter metadata for every initialized tensor.Source: Path instructions
222-249: 🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy liftFix
AttributeErroronNonevalues and accurately trackref_counterfor offloaded tensors.This implementation has several critical issues that will lead to crashes and memory leaks:
- Crash on
None:value._basewill raise anAttributeErrorifvalueisNone(which is permitted by thetorch.Tensor | Nonesignature and actively evaluated downstream).- Key mismatch memory leak:
ref_counteris incremented using the pre-offloadvalue. For caches likeDiskCachethat create a new meta tensor during offloading,value(CPU tensor) andoffloaded(Meta tensor) are different objects.__delitem__decrementsref_counter[offloaded], leaving the originalvaluepinned in memory forever and prematurely deleting the disk cache.- Double counting on updates:
ref_counteris unconditionally incremented at the top. If an existing key is updated in-place (torch.is_same_sizeis true), the reference count incorrectly inflates. Furthermore, old tensor references are not decremented when a key is completely overwritten.Refactor
__setitem__to safely check for_baseand to strictly manageref_counteragainst the actual storedoffloadedtensor object. As per path instructions, review cache utilities for no memory leaks.🔒️ Proposed fix
def __setitem__(self, key: Hashable, value: torch.Tensor | None): """ Update the offloaded and onloaded values if the key exists, otherwise offload the value and add it to the cache. """ # capture slice data - if value._base is not None: + if value is not None and getattr(value, "_base", None) is not None: base, view_index = index_from_view(value) self.view_index[key] = view_index value = base elif key in self.view_index: del self.view_index[key] - self.ref_counter[value] += 1 - # when onloading is disabled, parameters can be access and assigned directly if self.onloading_disabled: + old_value = self.offloaded_values.get(key, None) + if old_value is not None and old_value is not value: + self.ref_counter[old_value] -= 1 + if self.ref_counter[old_value] <= 0: + del self.ref_counter[old_value] self.offloaded_values[key] = value + if value is not None: + self.ref_counter[value] += 1 return # if the key already exists, update with the new value offloaded = self.offloaded_values.get(key, None) - if offloaded is not None and torch.is_same_size(offloaded, value): + if offloaded is not None and value is not None and torch.is_same_size(offloaded, value): self.update_offload(offloaded, value) onloaded = self.keep_onloaded_values.get(offloaded, None) if onloaded is not None and onloaded is not offloaded: onloaded.copy_(value) # if the key does not exist (or the value is None), offload the new value else: - self.offloaded_values[key] = self.offload(value) + if offloaded is not None: + self.ref_counter[offloaded] -= 1 + if self.ref_counter[offloaded] <= 0: + del self.ref_counter[offloaded] + + new_offloaded = self.offload(value) + self.offloaded_values[key] = new_offloaded + if new_offloaded is not None: + self.ref_counter[new_offloaded] += 1🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/compressed_tensors/offload/cache/base.py` around lines 222 - 249, Refactor __setitem__ to handle None before accessing _base, and manage ref_counter using the exact tensor object stored in offloaded_values after offload. Avoid incrementing counts for in-place updates, decrement the previous stored tensor when replacing a key, and keep __delitem__ consistent so every stored reference is released exactly once.Source: Path instructions
🧹 Nitpick comments (1)
src/compressed_tensors/offload/cache/disk.py (1)
77-79: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename
sliceto avoid shadowing the Python builtin.♻️ Proposed fix
- onloaded = file.get_slice(weight_info["weight_name"]) - slice = self.view_index.get(key, ...) - onloaded = onloaded[slice] # this materializes the tensor from `get_slice` + onloaded = file.get_slice(weight_info["weight_name"]) + view_slice = self.view_index.get(key, ...) + onloaded = onloaded[view_slice] # this materializes the tensor from `get_slice`🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/compressed_tensors/offload/cache/disk.py` around lines 77 - 79, Rename the local variable `slice` in the cache loading logic to a non-conflicting name, and update the subsequent `onloaded[...]` indexing expression to use it.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/compressed_tensors/offload/cache/cpu.py`:
- Around line 20-31: Add an early None return in onload for both
src/compressed_tensors/offload/cache/cpu.py lines 20-31 and
src/compressed_tensors/offload/cache/device.py lines 32-44, then rename the
slicing variable from slice to view_slice and update its use. Correct each
onload docstring so the tensor parameter is documented as offloaded rather than
duplicating key.
---
Outside diff comments:
In `@src/compressed_tensors/offload/cache/base.py`:
- Around line 124-129: Update from_mapping to populate the instance through
__setitem__ by assigning each mapping entry with instance[name] = tensor instead
of directly replacing offloaded_values. Preserve the existing offload behavior
while ensuring __setitem__ records view_index and ref_counter metadata for every
initialized tensor.
- Around line 222-249: Refactor __setitem__ to handle None before accessing
_base, and manage ref_counter using the exact tensor object stored in
offloaded_values after offload. Avoid incrementing counts for in-place updates,
decrement the previous stored tensor when replacing a key, and keep __delitem__
consistent so every stored reference is released exactly once.
---
Nitpick comments:
In `@src/compressed_tensors/offload/cache/disk.py`:
- Around line 77-79: Rename the local variable `slice` in the cache loading
logic to a non-conflicting name, and update the subsequent `onloaded[...]`
indexing expression to use it.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 51d715cc-0604-4d7d-b13d-d3e03bfec14b
📒 Files selected for processing (10)
src/compressed_tensors/offload/README.mdsrc/compressed_tensors/offload/cache/base.pysrc/compressed_tensors/offload/cache/cpu.pysrc/compressed_tensors/offload/cache/device.pysrc/compressed_tensors/offload/cache/disk.pysrc/compressed_tensors/offload/load.pysrc/compressed_tensors/offload/module.pysrc/compressed_tensors/offload/utils.pytests/test_offload/cache/helpers.pytests/test_offload/test_utils.py
🔗 Linked repositories identified
CodeRabbit considers these linked repositories for cross-repo context during reviews:
vllm-project/llm-compressor(manual)
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
kylesayrs
force-pushed
the
kylesayrs/offload-view
branch
from
August 12, 2026 16:58
365dd3a to
2f12bfd
Compare
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require one maintainer reviewWaiting for any of
This rule is failing.All PRs must have at least one approving review from a maintainer before merging.
|
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.
Purpose
Changes
view_indextoOffloadCachewhich tracks the views of each tensor relative to the base tensorref_counterwhich counts how many references there are to a base tensor via views. This can replaced by finalizers later if neededkeyargument toOffloadCache.onloadwhich allows lookup of the view associated with the onloaded tensorDisk Offloading Note
get_slicefollowed by an indexing only loads the necessary memory view from disk, not the entire tensorTesting