Simple tests

Author: S1ro1

tests/test_serialization.py

@@ -249,6 +249,31 @@ def test_get_torch_storage_size():
     assert get_torch_storage_size(torch.tensor([1, 2, 3, 4, 5], dtype=torch.float64)) == 5 * 8
     assert get_torch_storage_size(torch.tensor([1, 2, 3, 4, 5], dtype=torch.float16)) == 5 * 2
 
+@requires("torch")
+def test_get_torch_storage_size_dtensor():
+    # testing distributed sharded tensors isn't very easy, would need to subprocess call torchrun, so this should be good enough
+    import torch
+    import torch.distributed as dist
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.tensor import DTensor, Replicate
+
+    if dist.is_available() and not dist.is_initialized():
+        dist.init_process_group(
+            backend="gloo",
+            store=dist.HashStore(),
+            rank=0,
+            world_size=1,
+        )
+
+    mesh = init_device_mesh("cpu", (1,))
+    local = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float16)
+    dt = DTensor.from_local(local, mesh, [Replicate()])
+
+    assert get_torch_storage_size(dt) == 5 * 2
+
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
 
 @requires("torch")
 @pytest.mark.skipif(not is_wrapper_tensor_subclass_available(), reason="requires torch 2.1 or higher")