[Bugfix] Add synchronize to prevent possible data race (#6788)

Co-authored-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
This commit is contained in:
Tyler Michael Smith
2024-07-25 10:40:01 -07:00
committed by GitHub
co-authored by Lucas Wilkinson
parent 65b1f121c8
commit 95db75de64
+7
View File
@@ -243,6 +243,13 @@ class GroupCoordinator:
ca_comm = self.ca_comm
maybe_ca_context = nullcontext(
) if ca_comm is None else ca_comm.capture()
# ensure all initialization operations complete before attempting to
# capture the graph on another stream
curr_stream = torch.cuda.current_stream()
if curr_stream != stream:
stream.wait_stream(curr_stream)
with torch.cuda.stream(stream), maybe_ca_context:
# In graph mode, we have to be very careful about the collective
# operations. The current status is: