[FSDP2] Added missing event wait (for future) (pytorch#132568)

Nothing is actually wrong currently, but we should add this in case we land pytorch#127032 in the future.

Pull Request resolved: pytorch#132568
Approved by: https://github.com/weifengpy, https://github.com/Skylion007

torch/distributed/_composable/fsdp/_fsdp_state.py

@@ -290,7 +290,11 @@ def _root_post_backward_final_callback(self) -> None:
                     state._finalize_backward()
             if self._state_ctx.is_last_backward:
                 self._comm_ctx.post_forward_order.clear()
-                self._comm_ctx.reduce_scatter_state = None
+                if self._comm_ctx.reduce_scatter_state is not None:
+                    torch.cuda.current_stream().wait_event(
+                        self._comm_ctx.reduce_scatter_state.event
+                    )
+                    self._comm_ctx.reduce_scatter_state = None
             self._state_ctx.post_backward_final_callback_queued = False
 
     def _finalize_backward(self) -> None: