Add an option for TP-only AMAX reduction (#431)

Signed-off-by: Sangkug Lym <slym@nvidia.com>

transformer_engine/pytorch/fp8.py

@@ -326,24 +326,29 @@ class FP8GlobalStateManager:
             return None
 
         # Reduce AMAX in DP-domain at an interval.
+        # `NVTE_DP_AMAX_REDUCE_INTERVAL` should be set as an integer value larger than 0. If
+        # `NVTE_DP_AMAX_REDUCE_INTERVAL` is set to 0, AMAX is reduced only in TP domain.
         if cls.dp_amax_reduce_interval is None:
             cls.dp_amax_reduce_interval = int(os.getenv("NVTE_DP_AMAX_REDUCE_INTERVAL", "1"))
 
-        tp_amax_reduce = False
-        if forward:
-            if cls.dp_amax_reduce_forward_idx == 0:
-                reduce_group = fp8_meta["fp8_group"]
-            else:
-                tp_amax_reduce = True
-            cls.dp_amax_reduce_forward_idx = (
-                (cls.dp_amax_reduce_forward_idx + 1) % cls.dp_amax_reduce_interval)
+        if cls.dp_amax_reduce_interval == 0:
+            tp_amax_reduce = True
         else:
-            if cls.dp_amax_reduce_backward_idx == 0:
-                reduce_group = fp8_meta["fp8_group"]
+            tp_amax_reduce = False
+            if forward:
+                if cls.dp_amax_reduce_forward_idx == 0:
+                    reduce_group = fp8_meta["fp8_group"]
+                else:
+                    tp_amax_reduce = True
+                cls.dp_amax_reduce_forward_idx = (
+                    (cls.dp_amax_reduce_forward_idx + 1) % cls.dp_amax_reduce_interval)
             else:
-                tp_amax_reduce = True
-            cls.dp_amax_reduce_backward_idx = (
-                (cls.dp_amax_reduce_backward_idx + 1) % cls.dp_amax_reduce_interval)
+                if cls.dp_amax_reduce_backward_idx == 0:
+                    reduce_group = fp8_meta["fp8_group"]
+                else:
+                    tp_amax_reduce = True
+                cls.dp_amax_reduce_backward_idx = (
+                    (cls.dp_amax_reduce_backward_idx + 1) % cls.dp_amax_reduce_interval)
 
         if tp_amax_reduce:
             if tp_size > 1: