[dtensor] fix simplefsdp mixed-precision training bugs (#154975)

This is a follow-up on the previous dtensor redistribute PR: #150740, which enables SimpleFSDP's mixed-precision training.

In the most recent integration in TorchTitan: pytorch/torchtitan#1250, we found some discrepancies between SimpleFSDP's `fully_shard` and `replicate` modes when MPT is enabled. After debugging, I found the problem is in dtensor redistribute --`local_tensor` is taken out again from the original `input`. Thus, the dtensor used for communication has its original precision instead of using `forward_dtype`.

This PR fixes this issue and corrects previously added test cases.

After fixing the bug, the loss curves of `fully_shard` and `replicate` mode match perfectly.

![loss](https://github.com/user-attachments/assets/a8faddae-a476-48c0-a411-3fe04d2233bd)

Pull Request resolved: #154975
Approved by: https://github.com/tianyu-l

Author: ruisizhang123

test/distributed/tensor/test_redistribute.py

@@ -312,7 +312,9 @@ def test_shard_to_replicate_forward_backward_datatype_conversion(self):
                         backward_dtype=backward_dtype,
                     )
                 self.assertEqual(reshard_dtensor.size(), torch.Size(input_size))
-                self.assertEqual(expected_tensor, reshard_dtensor.to_local())
+                self.assertEqual(
+                    expected_tensor.to(forward_dtype), reshard_dtensor.to_local()
+                )
                 self.assertEqual(
                     comm_mode.get_comm_counts()[funcol.all_gather_into_tensor], 1
                 )

torch/distributed/tensor/_redistribute.py

@@ -318,7 +318,6 @@ def forward(  # type: ignore[override]
                 device_mesh, placements, tensor_meta=current_spec.tensor_meta
             )
 
-            local_tensor = input._local_tensor
             output = redistribute_local_tensor(
                 local_tensor, current_spec, target_spec, async_op=async_op
             )