[FSDP2] privateuse1 support fsdp2. (pytorch#139539)

We are looking forward to supporting FSDP2 with devices other than CUDA. Please give me some coding suggestions. Thank you very much.

Pull Request resolved: pytorch#139539
Approved by: https://github.com/kwen2501

torch/distributed/_composable/fsdp/_fsdp_param_group.py

@@ -51,8 +51,6 @@ class FSDPCommContext:
 
     def lazy_init(self, device: torch.device):
         self.device_handle = _get_device_handle(device.type)
-        if device.type not in ["cuda", "hpu"]:
-            raise RuntimeError("FSDP requires streams support")
         # Setting the all-gather/reduce-scatter streams to be higher priority
         # can help avoid some issues where their copies in/out are delayed and
         # block computation (this is different from high-pri NCCL streams)