[New] [Offload] [1/2] Disambiguate synchronous __setitem__/offload from update_offload

[kylesayrs]

Background

The ability to update offloaded values is something that can be done asynchronously across ranks.
The ability to create new offloaded values is something that must be done synchronously across ranks.

Currently on main, we do both through the same interface:

# this is a creation if the shapes don't match
# and an update if the shapes do match
module.weight = torch.nn.Parameter(new_value, ...)

# when updating, we assume that the shapes always match
def update_offload_parameter():
    ...
    # triggers update iff shapes match
    cache[name] = data

This is bad because it's very unclear to users when parameters are being updated vs created, which has serious implications for rank synchronization. It's better to make it very explicit: assignment is synchronous, while update_offload_parameter is always asynchronous.

Summary

• Always offload values when assigning parameters
  • Explicitly the key if it already exists: better cleanup for disk offloading
• Force update_offload_parameter to call update_offload explicitly
  • Make this function respect onloading_disabled
• Miscellaneous
  • Make tensors contiguous before writing to disk, as required by safetensors
  • Use torch.no_grad decorator for update_offload_parameter
  • Add offload_folder fixture for easy testing
  • Add test_distributed_compression_with_disk_offload for better coverage

Testing

• Update low-level unit tests to reflect new behavior of assignment
• https://github.com/neuralmagic/llm-compressor-testing/actions/runs/33812157889

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR changes offload update paths to validate stored tensors explicitly, delete existing cache entries before replacement, and handle missing disk-index entries by skipping updates. Tests were updated for local, GPU, and distributed update flows.

Changes

Offload update behavior

Layer / File(s)	Summary
OffloadCache setitem rework src/compressed_tensors/offload/cache/base.py	Existing keys are deleted before the new value is re-offloaded and stored.
update_offload_parameter rewrite src/compressed_tensors/offload/__init__.py	update_offload_parameter now runs under @torch.no_grad(), updates through cache.update_offload(...) for offloaded values, and only copies data in the non-cache path.
DiskCache update skip src/compressed_tensors/offload/cache/disk.py	update_offload now returns early when the offloaded tensor is absent from the index, after checking the incoming data shape/device state.
Local and GPU semantics tests tests/test_offload/test_interface.py, tests/test_offload/test_module.py, tests/test_offload/cache/test_disk.py	Adds coverage for offload-only updates, adjusts existing assertions, and updates overwrite behavior checks and a comment.
Distributed update paths tests/test_offload/cache/test_dist_cpu.py, tests/test_offload/cache/test_dist_disk.py, tests/test_offload/convert/test_from_accelerate.py	Distributed tests now read the current offloaded tensor and update through update_offload or update_offload_parameter instead of direct assignment or in-place mutation.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

• vllm-project/compressed-tensors#666: Changes DiskCache.update_offload validation and update behavior in the same file.
• vllm-project/compressed-tensors#709: Touches the same offload update paths and test coverage around them.
• vllm-project/compressed-tensors#767: Also modifies update_offload_parameter and OffloadCache.__setitem__.

Suggested reviewers: HDCharles, brian-dellabetta

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies the main change: separating synchronous setitem/offload behavior from update_offload. The release and series markers do not make it misleading.
Description check	✅ Passed	The description directly explains the synchronous assignment and asynchronous update changes, their synchronization implications, and the related tests.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch kylesayrs/synchronous-setitem

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @kylesayrs.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[coderabbitai]

Actionable comments posted: 1

🤖 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 `@tests/test_offload/test_interface.py`:
- Around line 98-111: The assertion in test_update_offload_parameter_only is
using value inequality on tensors, which can succeed for the wrong reason if
both tensors contain the same value. Update the checks around
offloaded_linear.weight and the local offload variable to assert they are
distinct by identity or by comparing device/state, using the existing test
helpers disable_offloading and disable_onloading to verify the onloaded and
offloaded tensors are not the same object.

🪄 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: 219c97bf-7825-4727-9f61-43365bbb801e

📥 Commits

Reviewing files that changed from the base of the PR and between a4ee36a and f05c0c9.

📒 Files selected for processing (5)

• src/compressed_tensors/offload/__init__.py
• src/compressed_tensors/offload/cache/base.py
• tests/test_offload/cache/test_disk.py
• tests/test_offload/test_interface.py
• tests/test_offload/test_module.py

🔗 Linked repositories identified

CodeRabbit considers these linked repositories for cross-repo context during reviews:

• vllm-project/llm-compressor (manual)

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# 1. Check device constants in test_interface.py
rg -n 'ONLOAD_DEVICE|OFFLOAD_DEVICE' tests/test_offload/test_interface.py | head -10

# 2. Check CPUCache.update_offload implementation
ast-grep outline src/compressed_tensors/offload/cache/cpu.py --match update_offload --view expanded

# 3. If outline doesn't show the body, read the method directly
rg -n -A 15 'def update_offload' src/compressed_tensors/offload/cache/cpu.py

Repository: vllm-project/compressed-tensors

Length of output: 1085

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== tests/test_offload/test_interface.py =="
sed -n '1,180p' tests/test_offload/test_interface.py

echo
echo "== src/compressed_tensors/offload/cache/cpu.py =="
sed -n '1,220p' src/compressed_tensors/offload/cache/cpu.py

echo
echo "== src/compressed_tensors/offload/interface.py =="
sed -n '1,260p' src/compressed_tensors/offload/interface.py

Repository: vllm-project/compressed-tensors

Length of output: 7714

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

rg -n -A 80 'def update_offload_parameter|keep_onloaded_values|disable_offloading|disable_onloading' src/compressed_tensors/offload tests/test_offload

Repository: vllm-project/compressed-tensors

Length of output: 50390

---

Use identity or device comparison here
!= compares tensor values, so this can pass/fail for the wrong reason once both tensors contain 1. Use is not (or compare .device) to assert the onloaded and offloaded tensors are distinct.

🤖 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 `@tests/test_offload/test_interface.py` around lines 98 - 111, The assertion in
test_update_offload_parameter_only is using value inequality on tensors, which
can succeed for the wrong reason if both tensors contain the same value. Update
the checks around offloaded_linear.weight and the local offload variable to
assert they are distinct by identity or by comparing device/state, using the
existing test helpers disable_offloading and disable_onloading to verify the
onloaded and offloaded tensors are not the same object.

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require one maintainer review	👀 reviews

🔴 Require one maintainer review

Waiting for any of

• approved-reviews-by=HDCharles
• approved-reviews-by=brian-dellabetta
• approved-reviews-by=dsikka
• approved-reviews-by=kylesayrs

This rule is failing.

All PRs must have at least one approving review from a maintainer before merging.

• any of:
  • approved-reviews-by=HDCharles
  • approved-reviews-by=brian-dellabetta
  • approved-reviews-by=dsikka
  • approved-reviews-by=kylesayrs
• #changes-requested-reviews-by = 0