This PR adds mixed precision training support for SimpleFSDP, togetherwith this change from PyTorch:pytorch/pytorch#150740.**(This is a placeholder for now, and shall be merged only after thePyTorch pr is landed.)**#### Convergence on debug modelThe loss curves are all benchmarked with `training.seed = 42` on 4 GPUs.- Fully Shard Mode (['dp_shard'], [4]): FSDP2 and SimpleFSDP's lossesare perfectly matched<img width="2234" alt="Screenshot 2025-04-05 at 4 06 04 PM"src="https://github.com/user-attachments/assets/e0ec66f2-948f-43e2-8d88-20243fbaccfc"/>- Hybrid Shard Mode (['dp_replicate', 'dp_shard'], [2, 2]): FSDP2 andSimpleFSDP's losses are similar<img width="977" alt="Screenshot 2025-04-07 at 6 15 04 PM"src="https://github.com/user-attachments/assets/81c14637-e0a7-42aa-9b2e-c7092a287ea0"/>The end-to-end SimpleFSDP mixed precision training integration has beenproved to work properly in the PR from this fork:tianyu-l/pytorch_intern24#20.

…ble SimpleFSDP mixed precision training (pytorch#150740)As titled, this pr adds additional `forward_dtype` and `backward_dtype` conversion in DTensor `redistribute` API to enable SimpleFSDP's mixed precision training.In this forward pass, the DTensor can be configured to be cast to `forward_dtype`; in the backward pass, the DTensor can be configured to be cast to `backward_dtype`.1. **Correctness**: The end-to-end SimpleFSDP mixed precision training integration has been proved to work properly in the PR from this fork:tianyu-l#20. We are now migrating the code to official PyTorch DTensor.2. **Example Usage**: There is an example in TorchTian's SimpleFSDP implementation:pytorch/torchtitan#1060.In the example below, a DTensor `x` is all-gather'ed along the `self.compute_placements`, with datatype cast to `self.param_dtype`. In the backward pass, additionally, the computed gradients are reduce-scatter'ed along the `self.grad_placements`, with datatype cast to `self.reduce_dtype`.```pythonoutput = x.redistribute(        placements=self.compute_placements,        forward_dtype=self.param_dtype,        backward_dtype=self.reduce_dtype,).to_local(grad_placements=self.grad_placements)```Under the hood, in `class Redistribute(torch.autograd.Function):`, the `forward` function first takes `x`'s local tensor, convert it to `forward_dtype`, before all-gather `x`.The `backward` function take `grad_output` and convert it to `backward_dtype`, before reduce-scatter `grad_output`.Pull Requestresolved:pytorch#150740Approved by:https://github.com/tianyu-l

…ble SimpleFSDP mixed precision training (pytorch#150740)As titled, this pr adds additional `forward_dtype` and `backward_dtype` conversion in DTensor `redistribute` API to enable SimpleFSDP's mixed precision training.In this forward pass, the DTensor can be configured to be cast to `forward_dtype`; in the backward pass, the DTensor can be configured to be cast to `backward_dtype`.1. **Correctness**: The end-to-end SimpleFSDP mixed precision training integration has been proved to work properly in the PR from this fork:tianyu-l#20. We are now migrating the code to official PyTorch DTensor.2. **Example Usage**: There is an example in TorchTian's SimpleFSDP implementation:pytorch/torchtitan#1060.In the example below, a DTensor `x` is all-gather'ed along the `self.compute_placements`, with datatype cast to `self.param_dtype`. In the backward pass, additionally, the computed gradients are reduce-scatter'ed along the `self.grad_placements`, with datatype cast to `self.reduce_dtype`.```pythonoutput = x.redistribute(        placements=self.compute_placements,        forward_dtype=self.param_dtype,        backward_dtype=self.reduce_dtype,).to_local(grad_placements=self.grad_placements)```Under the hood, in `class Redistribute(torch.autograd.Function):`, the `forward` function first takes `x`'s local tensor, convert it to `forward_dtype`, before all-gather `x`.The `backward` function take `grad_output` and convert it to `backward_dtype`, before reduce-scatter `grad_output`.Pull Requestresolved:pytorch#150740Approved by:https://github.com/tianyu-l

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 Requestresolved:#154975Approved by:https://github.com/tianyu-l

This is a follow-up on the previous dtensor redistribute PR:pytorch#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 Requestresolved:pytorch#154975Approved by:https://github.com/tianyu-l

This is a follow-up on the previous dtensor redistribute PR:pytorch#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 Requestresolved:pytorch#154975Approved by:https://github.com/tianyu-l

This is a follow-up on the previous dtensor redistribute PR:pytorch#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 Requestresolved:pytorch#154975Approved by:https://github.com/tianyu-l