[Bug] Disk Offloading and Deletion of Shared Tensors

[kylesayrs]

Background

Deleting shared parameters means that the disk file is deleted twice, leading to an error

def test_shared_tensors_delete(tmp_path):
    offload_dir = tmp_path / "offload_dir"
    os.mkdir(offload_dir)

    A = torch.nn.Linear(3, 5, bias=False)
    B = torch.nn.Linear(3, 5, bias=False)

    offload_module(A, "cuda", "disk", offload_dir=str(offload_dir))
    with disable_onloading():
        B.weight = A.weight
    offload_module(B, "cuda", "disk", offload_dir=str(offload_dir))
    assert len(os.listdir(offload_dir)) == 1

    delattr(A, "weight")
    assert len(os.listdir(offload_dir)) == 1
    delattr(B, "weight")
    assert len(os.listdir(offload_dir)) == 0

Proposed Changes

• Design a solution to allow for deletion of shared tensors when using the disk cache

[Yatimai]

Hi @kylesayrs, I'd like to take a stab at this.

I traced the deletion path and the crash comes from the from_accelerate flow with tied weights (e.g. lm_head.weight ↔ embed_tokens.weight). Both modules end up with the same meta tensor T in their respective DiskCache instance, pointing to the same file via DiskCache.index[T]. First __delitem__ deletes the file and removes T from index. Second __delitem__ hits a KeyError on self.index[offloaded] at disk.py:118.

Proposed approach: class-level refcount on DiskCache, same pattern as index.

• _file_refcount: dict[torch.Tensor, int] = dict() on DiskCache
• offload(), increment on both the new-tensor and meta-reuse branches
• __delitem__(), decrement; only delete file + index entry when count reaches 0; always call super().__delitem__() for offloaded_values cleanup
• DistributedDiskCache, same logic in offload() (non-rank-0 branch) and __delitem__()
• No changes to update_offload or create_checkpoint_symlink