Add dtensor support for TE optimizers (#1171)

add dtensor support for te optimizers
Signed-off-by: jasonwan <jasonwan@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

transformer_engine/pytorch/optimizers/multi_tensor_apply.py

@@ -3,6 +3,7 @@
 # See LICENSE for license information.
 
 """Multi-tensor apply entry."""
+from torch.distributed._tensor import DTensor
 
 
 class MultiTensorApply:  # pylint: disable=too-few-public-methods
@@ -12,6 +13,11 @@ class MultiTensorApply:  # pylint: disable=too-few-public-methods
         self.chunk_size = chunk_size
 
     def __call__(self, op, noop_flag_buffer, tensor_lists, *args):
+        for i, ts in enumerate(tensor_lists):
+            for j, t in enumerate(ts):
+                if isinstance(t, DTensor):
+                    tensor_lists[i][j] = t._local_tensor
+
         return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args)