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

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

transformer_engine/pytorch/fp8.py

@@ -326,9 +326,14 @@ 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"))
 
+        if cls.dp_amax_reduce_interval == 0:
+            tp_amax_reduce = True
+        else:
             tp_amax_reduce = False
             if forward:
                 if cls.dp_amax_reduce_forward_idx == 0: