trunk/31527a43dbf6adf4df2e6eb5c7b38094fec6b6f6: Widen FSDP param group stream type annotations to torch.Stream (#196563)
- PyTorch: 784 events in the last 90 days
- PyTorch: 774th Release in the last 90 days
- Previous: earlier the same day · trunk/6e700ca9014c0c7ba112f66d02a2e81b94cc1f20: Handle non-UTF-8 subprocess output in minifier isolate_fails (#196620)
What happened
Summary Widen the type annotations of _all_reduce_hook_stream and all_reduce_stream in FSDPParamGroup from torch.cuda.Stream to torch.Stream , allowing out-of-tree accelerator backends to pass their own Stream subtypes through FSDP's all-reduce hook path. Changes torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py : L208 self._all_reduce_hook_stream: torch.cuda.Stream | None → torch.Stream | None torch/distribu…
Summary assembled by rule from the sources below