Broadcast tensor shape in distributed offload caches

[kylesayrs]

Stacked on top of #856. Review/merge that one first.

Purpose

The distributed CPU and disk offload caches only reconciled dtype divergence between ranks; the non-source (meta-init) rank kept its own tensor shape. For tied/multimodal models, transformers can meta-initialize a parameter with a different shape than the checkpoint tensor materialized on the source rank.

In the CPU cache this makes Tensor.set_() lay the non-source rank's larger layout over the source's smaller shared storage:

RuntimeError: setStorage: sizes [128], strides [1], storage offset 0, and itemsize 4
requiring a storage size of 512 are out of bounds for storage of size 256

Changes

• Broadcast the source tensor's shape alongside its dtype in DistributedCPUCache.offload and DistributedDiskCache.offload.
• Rebuild the non-source tensor to match both dtype and shape before pointing it at the shared storage (CPU) / recording the meta placeholder (disk).
• Add test_load_dist_no_missing_keys, covering multimodal, tied, and untied models under load_offloaded_model in a 2-rank (world_size=2) setup.

Testing

test_load_dist_no_missing_keys passes on 2 GPUs across all four models listed in #856. Without this change, the multimodal/tied gemma-4-1B-0.8B-tiny case fails with the setStorage error above.

Note: full disk offload can't be exercised in CI without an accelerate upgrade (ValueError: You are trying to offload the whole model to the disk), so the disk cache carries the analogous fix for parity but is not covered by a new test.

https://github.com/neuralmagic/llm-compressor-testing/actions/runs/33012761751

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 5a094839-23ee-465f-ae94-916f1af894f6

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

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]

Merge Protections

🟢 Merge protection satisfied — ready to merge.

Show 1 satisfied protection

🟢 Require one maintainer review

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

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

[brian-dellabetta]

LGTM, comment is helpful!

[mergify]

The quality checks have failed. Please run make style and make quality under
the root directory to adddress the lint failures. You will need to install the
dev optional install to get the required linting packages.

[mergify]

The quality checks have failed. Please run make style and make quality under
the root directory to adddress the lint failures. You will need to install the
dev optional install to get the required linting packages.