[XLA] Implement XLAShardedTensor.to_local()

Author: Hoomaaan

test/neuron/run_tests.sh

@@ -257,6 +257,7 @@ function run_xla_op_tests3 {
   run_test_multi_device "$_TEST_DIR/spmd/test_dtensor_convert_mesh.py"
   run_test_multi_device "$_TEST_DIR/spmd/test_xla_dtensor_spec_conv.py"
   run_test_multi_device "$_TEST_DIR/spmd/test_xla_dtensor_from_local.py"
+  run_test_multi_device "$_TEST_DIR/spmd/test_xla_dtensor_to_local.py"
   run_test "$_TEST_DIR/spmd/test_xla_auto_sharding.py"
   #run_test "$_TEST_DIR/spmd/test_spmd_parameter_wrapping.py"
   run_test "$_TEST_DIR/spmd/test_train_spmd_linear_model.py"

test/run_tests.sh

@@ -256,6 +256,7 @@ function run_xla_op_tests3 {
   run_test_multi_devices "$_TEST_DIR/spmd/test_dtensor_convert_mesh.py"
   run_test_multi_devices "$_TEST_DIR/spmd/test_xla_dtensor_spec_conversion.py"
   run_test_multi_devices "$_TEST_DIR/spmd/test_xla_dtensor_from_local.py"
+  run_test_multi_devices "$_TEST_DIR/spmd/test_xla_dtensor_to_local.py"
   run_test "$_TEST_DIR/spmd/test_xla_auto_sharding.py"
   run_test "$_TEST_DIR/spmd/test_spmd_parameter_wrapping.py"
   run_test "$_TEST_DIR/spmd/test_mp_input_sharding.py"

test/spmd/test_xla_dtensor_to_local.py

@@ -0,0 +1,75 @@
+import sys
+import unittest
+import torch
+import numpy as np
+
+from torch.distributed.tensor import DeviceMesh
+from torch.distributed._tensor import DTensor
+from torch.distributed.tensor.placement_types import Replicate, Shard
+import torch_xla
+import torch_xla.runtime as xr
+import torch_xla.core.xla_model as xm
+from torch_xla.distributed.spmd.xla_sharded_tensor import XLAShardedTensor
+import test_xla_sharding_base
+
+
+class DTensorXLAFromLocalConversionTest(test_xla_sharding_base.XlaShardingTest):
+    """
+    Test suite for the automatic conversion of regular tensors to XLAShardedTensor
+    in DTensor.from_local() when using XLA device mesh.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+
+    def test_to_local(self):
+        from torch.distributed.tensor import distribute_tensor
+        world_size = xr.global_runtime_device_count()
+        mesh = DeviceMesh("xla", list(range(world_size)))
+        
+        big_tensor = torch.randn(100000, 88)
+        sharded_tensor = XLAShardedTensor(big_tensor, mesh, [Shard(0)])
+  
+        local_tensor = sharded_tensor.to_local()
+
+        # Verify the shapes are the same
+        self.assertEqual(local_tensor.shape, big_tensor.shape)
+
+        # Check the value of the tensor
+        torch.testing.assert_close(local_tensor, big_tensor, check_device=False)
+
+    def test_to_local_requires_grad(self):
+        """Test that gradients flow correctly through to_local()."""
+        # Create a tensor with requires_grad=True
+        world_size = xr.global_runtime_device_count()
+        mesh = DeviceMesh("xla", list(range(world_size)))
+
+        tensor = torch.randn(100_000, 88, requires_grad=True)
+        
+        # Create XLAShardedTensor
+        sharded_tensor = XLAShardedTensor(tensor, mesh, [Shard(0)])
+        
+        # Verify requires_grad is set
+        self.assertTrue(sharded_tensor.requires_grad)
+        
+        res = sharded_tensor.sum()
+        res.backward()
+
+        # Verify grad are calculated
+        self.assertTrue(sharded_tensor.grad is not None)
+
+        # Call to local function
+        local_tensor = sharded_tensor.to_local()
+        
+        # Verify requires_grad is preserved
+        self.assertTrue(local_tensor.requires_grad)
+        
+        # All gradients should be 1.0 since we did a sum()
+        self.assertTrue(torch.allclose(local_tensor.grad, torch.ones_like(tensor)))
+        
+        print("Gradient flow test successful")
+
+if __name__ == "__main__":
+    result = unittest.main(exit=False)
+    sys.exit(0 if result.result.wasSuccessful() else 1)

test/tpu/run_tests.sh

@@ -63,6 +63,7 @@ run_test "$_TEST_DIR/spmd/test_fsdp_v2.py"
 run_test "$_TEST_DIR/spmd/test_dtensor_convert_mesh.py"
 run_test "$_TEST_DIR/spmd/test_xla_dtensor_spec_conversion.py"
 run_test "$_TEST_DIR/spmd/test_xla_dtensor_from_local.py"
+run_test "$_TEST_DIR/spmd/test_xla_dtensor_to_local.py"
 run_test "$_TEST_DIR/test_gradient_accumulation.py"
 XLA_EXPERIMENTAL=nonzero:masked_select:nms run_test "$_TEST_DIR/ds/test_dynamic_shape_models.py" -v
 run_test "$_TEST_DIR/test_autocast.py"

torch_xla/distributed/spmd/xla_sharded_tensor.py

@@ -114,7 +114,7 @@ def __new__(cls,
         dtype=elem.dtype,
         layout=elem.layout,
         device=elem.device,
-        requires_grad=kwargs.get("requires_grad", False))
+        requires_grad=kwargs.get("requires_grad", elem.requires_grad))
     r.global_tensor = elem.detach() if r.requires_grad else elem
 
     # Initialize mesh, partition, and spec information
@@ -150,6 +150,29 @@ def load_local_shards_(self, shards: List[XLAShard]):
     # Invalidate cached spec since the global_tensor data has changed
     self.invalidate_spec_cache()
 
+  def to_local(self):
+    """
+    Returns the local representation of the XLAShardedTensor.
+    
+    This method returns the global tensor representation, which contains
+    the combined data across all devices. The returned tensor is on the
+    same device as the original XLAShardedTensor. The returned tensor
+    will have the same requires_grad value as the XLAShardedTensor.
+    If the original tensor has gradients, those will be preserved.
+    
+    Returns:
+        torch.Tensor: The global tensor representation with appropriate requires_grad setting.
+    """
+
+    # Create a new tensor with the same values of global_tensor
+    result = self.global_tensor.clone()
+    # Since global tensor is detached, add requires_grad and grad values back to the local tensor
+    if self.requires_grad:
+      result.requires_grad = self.requires_grad
+      result.grad = self.grad
+
+    return result
+
   @property
   def sharding_spec(self):
     return torch_xla._XLAC._get_xla_sharding_spec(self.global_tensor)