Merge branch 'main' into sanshang/fix_pg_init

facebookresearch · Nov 1, 2024 · 0ef5713 · 0ef5713
2 parents 61009fb + e13c79f
commit 0ef5713
Show file tree

Hide file tree

Showing 5 changed files with 218 additions and 85 deletions.
diff --git a/et_replay/tools/comm_replay.py b/et_replay/tools/comm_replay.py
@@ -804,10 +804,16 @@ def runComms(
             if collName in self.backendFuncs.collectiveFunc.keys():
                 # record wait_obj_key for wait ops
                 if curComm.req is not None and curComm.pgId is not None:
+                    if isinstance(curComm.req, list):
+                        seq_id = curComm.req[0]
+                        is_p2p_op = curComm.req[1]
+                    else:
+                        seq_id = curComm.req
+                        is_p2p_op = False
                     self.collectiveArgs.wait_obj_key = (
                         curComm.pgId,
-                        curComm.req[0],
-                        curComm.req[1],
+                        seq_id,
+                        is_p2p_op,
                     )
                 else:
                     self.collectiveArgs.wait_obj_key = None