fix error in #701 (#704)

Summary: I mistakenly landed a debug version of #701, fixing. Test Plan: CI Reviewers: Subscribers: Tasks: Tags:
pytorch · Nov 27, 2024 · 87e2c09 · 87e2c09
1 parent bcebc0d
commit 87e2c09
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchtitan/parallelisms/parallelize_llama.py b/torchtitan/parallelisms/parallelize_llama.py
@@ -238,8 +238,8 @@ def apply_tp(
     torch.ops._c10d_functional.reduce_scatter_tensor.default,
     # for low precision training, it's useful to always save
     # the result of max(abs(tensor))
-    # torch.ops.aten.abs.default,
-    # torch.ops.aten.max.default,
+    torch.ops.aten.abs.default,
+    torch.ops.aten.max.default,
 }