Address review comments and fix style issues

facebookresearch · Jun 27, 2024 · f1dd62f · f1dd62f
1 parent 94ffee3
commit f1dd62f
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 6 deletions.
diff --git a/et_replay/comm/commsTraceParser.py b/et_replay/comm/commsTraceParser.py
@@ -239,7 +239,9 @@ def _parseExecutionTrace(
                 ranks = pg["ranks"]
                 groupCnt = pg["group_count"]
                 pgRanksMap[pgId] = (
-                    ranks if len(ranks) > 0 else list(range(pg["group_size"]))
+                    ranks
+                    if len(ranks) > 0
+                    else list(range(pg["group_size"]))
                     # rank list is empty when all ranks are in a pg
                 )
                 if ET_BACKENDID:

diff --git a/et_replay/comm/comms_utils.py b/et_replay/comm/comms_utils.py
@@ -790,7 +790,9 @@ def __init__(self, args: Namespace) -> None:
         self.quant_a2a_embedding_dim = args.quant_a2a_embedding_dim
         self.quant_threshold = args.quant_threshold
         self.dcheck = args.c
-        self.groupRanks = {}  # record what ranks each process group will work on {pg_id, ranks}
+        self.groupRanks = (
+            {}
+        )  # record what ranks each process group will work on {pg_id, ranks}
         self.use_ext_dist = args.use_ext_dist
         self.size_from_trace = False
         self.init_method = args.init_method

diff --git a/et_replay/tools/et_replay.py b/et_replay/tools/et_replay.py
@@ -608,11 +608,11 @@ def allocate_tensors(self):
                             self.tensors_mapping[
                                 (node.id, tuple(node.inputs[2][:5]), True)
                             ]
-                        ][i] = i * nnz
+                        ][i] = (i * nnz)
                     else:
                         self.tensor_registry_permanent[
                             self.tensors_mapping[(node.id, tuple(node.inputs[2]), True)]
-                        ][i] = i * nnz
+                        ][i] = (i * nnz)
             ######
 
     def build_func(self, node):
@@ -1183,7 +1183,8 @@ def run_op(self, node, iter):
                 not in self.unchangeable_intermediate_tensors
             ):
                 if (
-                    self.tensors_mapping[(node.id, t_id, False)] not in self.instantiate
+                    self.tensors_mapping[(node.id, t_id, False)]
+                    not in self.instantiate
                     # and self.tensors_mapping[(node.id, t_id, False)]
                     # not in self.tensor_registry
                 ):

diff --git a/et_replay/utils.py b/et_replay/utils.py
@@ -6,7 +6,7 @@
 import uuid
 from typing import Any, Dict
 
-from et_replay.execution_trace import ExecutionTrace
+from et_replay import ExecutionTrace
 
 
 def get_tmp_trace_filename() -> str: