update

pytorch · May 10, 2024 · f9b969c · f9b969c
1 parent 90880a0
commit f9b969c
Showing 1 changed file with 1 addition and 6 deletions.
diff --git a/torch_xla/distributed/zero_redundancy_optimizer.py b/torch_xla/distributed/zero_redundancy_optimizer.py
@@ -419,8 +419,7 @@ def allgather_weights_and_update_full_parameter(self):
         if param.grad is not None or (self.use_grad_acc_hook and
                                       hasattr(shard, 'main_grad')):
           shard_data = shard.data
-          if not self.higher_cc_precision:
-            shard_data = shard_data.to(dtype=param.dtype)
+          shard_data = shard_data.to(dtype=param.dtype)
           if self.coalesce_cc_all_gather:
             sharded_data.append(shard_data)
           else:
@@ -430,8 +429,6 @@ def allgather_weights_and_update_full_parameter(self):
                 pin_layout=self.pin_layout,
                 groups=self.sharding_groups,
             )
-            if padded_param.dtype != param.dtype:
-              padded_param = padded_param.to(dtype=param.dtype)
             param.data.copy_(padded_param.data[:param.size(0)])
 
     if self.coalesce_cc_all_gather:
@@ -450,8 +447,6 @@ def allgather_weights_and_update_full_parameter(self):
           if param.grad is not None or (self.use_grad_acc_hook and
                                         hasattr(shard, 'main_grad')):
             padded_param = padded_params[index]
-            if padded_param.dtype != param.dtype:
-              padded_param = padded_params[index].to(dtype=param.dtype)
             param.data.copy_(padded_param.data[:param.size(0)])
             index += 1