fix the failing CI by reverting to default requires_grad

Author: Hoomaaan

test/spmd/test_xla_dtensor_to_local.py

@@ -44,7 +44,7 @@ def test_to_local_requires_grad(self):
     tensor = torch.randn(100_000, 88, requires_grad=True)
 
     # Create XLAShardedTensor
-    sharded_tensor = XLAShardedTensor(tensor, mesh, [Shard(0)])
+    sharded_tensor = XLAShardedTensor(tensor, mesh, [Shard(0)], requires_grad=tensor.requires_grad)
 
     # Verify requires_grad is set
     self.assertTrue(sharded_tensor.requires_grad)
@@ -70,7 +70,7 @@ def test_to_local_grad_independence(self):
     mesh = DeviceMesh("xla", list(range(world_size)))
 
     tensor = torch.randn(100_000, 88, requires_grad=True)
-    sharded_tensor = XLAShardedTensor(tensor, mesh, [Shard(0)])
+    sharded_tensor = XLAShardedTensor(tensor, mesh, [Shard(0)], requires_grad=tensor.requires_grad)
 
     # Create gradients
     res = sharded_tensor.sum()
@@ -95,7 +95,7 @@ def test_to_local_grad_none_handling(self):
     mesh = DeviceMesh("xla", list(range(world_size)))
 
     tensor = torch.randn(100_000, 88, requires_grad=True)
-    sharded_tensor = XLAShardedTensor(tensor, mesh, [Shard(0)])
+    sharded_tensor = XLAShardedTensor(tensor, mesh, [Shard(0)], requires_grad=tensor.requires_grad)
 
     # Don't do backward pass, so grad remains None
     self.assertIsNone(sharded_tensor.grad)

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", elem.requires_grad))
+        requires_grad=kwargs.get("requires_grad", False))
     r.global_tensor = elem.detach() if r.requires_grad else elem
 
     # Initialize mesh, partition, and spec information
@@ -164,14 +164,20 @@ def to_local(self):
         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.clone() if self.grad is not None else None
 
-    return result
+    if not self.requires_grad:
+      #  When requires_grad is False, global_tensor is the original tensor
+      return self.global_tensor
+    else:
+      # When requires_grad is True, global_tensor is detached
+      # 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.clone() if self.grad is not None else None
+
+      return result
 
   @property
   def sharding_spec(self):