Broadcast tensor shape in distributed offload caches

src/compressed_tensors/offload/cache/dist_cpu.py

@@ -6,7 +6,7 @@
 from compressed_tensors.distributed import get_source_rank, is_source_process
 from compressed_tensors.offload.cache.cpu import CPUCache
 from compressed_tensors.offload.cache.utils import catch_cpu_mem_error
-from compressed_tensors.offload.utils import send_tensors, to_empty
+from compressed_tensors.offload.utils import send_tensors, to_tensor
 
 
 class DistributedCPUCache(CPUCache):
@@ -19,11 +19,6 @@ def offload(self, tensor: torch.Tensor | None) -> torch.Tensor | None:
         """
         Synchronously create shared cpu memory for offload.
 
-        The dtype of ``tensor`` on non-source ranks cannot be trusted because
-        transformers may initialize buffers (e.g. ``inv_freq``) with a
-        different dtype than the checkpoint value on the source rank. See
-        https://github.com/huggingface/transformers/pull/47486
-
         :param tensor: tensor on any device
         :return: cpu tensor whose data is located in shared memory
         """
@@ -37,18 +32,31 @@ def offload(self, tensor: torch.Tensor | None) -> torch.Tensor | None:
             # create shared memory cpu tensor
             tensor = super().offload(tensor).share_memory_()
             handle, filename, nbytes = tensor.untyped_storage()._share_filename_cpu_()
-            broadcast_obj = [handle, filename, nbytes, tensor.dtype]
+            broadcast_obj = [handle, filename, nbytes, tensor.dtype, tensor.shape]
         else:
-            broadcast_obj = [None, None, None, None]
+            broadcast_obj = [None, None, None, None, None]
 
         # receive shared memory file handle
         dist.broadcast_object_list(broadcast_obj, src=get_source_rank())
 
         if not is_source_process():
+            src_shape = broadcast_obj.pop(4)
             src_dtype = broadcast_obj.pop(3)
 
-            if tensor.device.type == "meta" or tensor.dtype != src_dtype:
-                tensor = to_empty(tensor, device=self.offload_device, dtype=src_dtype)
+            # transformers may init params/buffers on non-source (meta) ranks with a
+            # different dtype or shape than the checkpoint (e.g. `inv_freq`, or
+            # tied/multimodal weights), so rebuild from the source's dtype and shape
+            # before pointing at the shared storage. See
+            # https://github.com/huggingface/transformers/pull/47486
+            if (
+                tensor.is_meta
+                or tensor.dtype != src_dtype
+                or tensor.shape != src_shape
+            ):
+                empty = torch.empty(
+                    src_shape, dtype=src_dtype, device=self.offload_device
+                )
+                tensor = to_tensor(empty, tensor)
             else:
                 tensor = send_tensors(tensor, device=self.offload_device)
 

src/compressed_tensors/offload/cache/dist_disk.py

@@ -5,7 +5,7 @@
 import torch.distributed as dist
 from compressed_tensors.distributed import get_source_rank, is_source_process
 from compressed_tensors.offload.cache.disk import DiskCache
-from compressed_tensors.offload.utils import send_tensors, to_empty
+from compressed_tensors.offload.utils import send_tensors, to_tensor
 
 
 class DistributedDiskCache(DiskCache):
@@ -18,11 +18,6 @@ def offload(self, tensor: torch.Tensor | None) -> torch.Tensor | None:
         """
         Synchronously write tensor data to disk.
 
-        The dtype of ``tensor`` on non-source ranks cannot be trusted because
-        transformers may initialize buffers (e.g. ``inv_freq``) with a
-        different dtype than the checkpoint value on the source rank. See
-        https://github.com/huggingface/transformers/pull/47486
-
         :param tensor: tensor on any device
         :return: meta tensor representing disk offloaded parameter
         """
@@ -36,17 +31,26 @@ def offload(self, tensor: torch.Tensor | None) -> torch.Tensor | None:
                 self.index[offloaded]["safetensors_file"],
                 self.index[offloaded]["weight_name"],
                 self.index[offloaded]["dtype"],
+                offloaded.shape,
             ]
         else:
             offloaded = send_tensors(tensor, device="meta")
-            broadcast_obj = [None, None, None]
+            broadcast_obj = [None, None, None, None]
 
         dist.broadcast_object_list(broadcast_obj, src=get_source_rank())
 
         if not is_source_process():
             src_dtype = getattr(torch, broadcast_obj[2])
-            if offloaded.dtype != src_dtype:
-                offloaded = to_empty(offloaded, device="meta", dtype=src_dtype)
+            src_shape = broadcast_obj[3]
+
+            # transformers may init params/buffers on non-source (meta) ranks with a
+            # different dtype or shape than the checkpoint (e.g. `inv_freq`, or
+            # tied/multimodal weights), so rebuild the meta tensor to match the
+            # source. See https://github.com/huggingface/transformers/pull/47486
+            if offloaded.dtype != src_dtype or offloaded.shape != src_shape:
+                empty = torch.empty(src_shape, dtype=src_dtype, device="meta")
+                offloaded = to_tensor(empty, offloaded)
+
             self.index[offloaded] = {
                 "safetensors_file": broadcast_obj[0],
                 "weight_name": broadcast_obj[1],

tests/test_offload/test_load.py

@@ -236,6 +236,14 @@ def test_load_no_missing_keys(model_id, model_class):
     _assert_load_no_missing_keys(model_id, model_class, device_map="cpu")
 
 
+@pytest.mark.integration
+@requires_gpu(2)
+@torchrun(world_size=2, init_dist=True)
+def test_load_dist_no_missing_keys():
+    for model_id, model_class in LOAD_NO_MISSING_KEYS_PARAMETERS:
+        _assert_load_no_missing_keys(model_id, model_class, device_map="cpu")
+
+
 @pytest.mark.integration
 @requires_gpu(2)
 @torchrun(world_size=2, init_dist=True)