minor changes in float8_experimental

Reviewed By: vkuzo

Differential Revision: D56867907

fbshipit-source-id: 27e61d7c14d2d406c19bfe728693ef988befa2e8

Author: y-sq

float8_experimental/float8_linear.py

@@ -183,6 +183,10 @@ def __init__(self, *args, **kwargs):
 
         # pre_forward and post_forward are currently broken with FSDP
         # and torch.compile, this option can disable them
+        # Note that when using `config.enable_pre_and_post_forward = False`,
+        # it's recommended to also set `config.enable_amax_init = False`.
+        # Otherwise, the amax buffer would never be marked as initialized and
+        # would be initialized in every iteration.
         self.enable_pre_and_post_forward = config.enable_pre_and_post_forward
 
     def register_always_float32_buffer(

float8_experimental/float8_utils.py

@@ -119,9 +119,9 @@ def fp8_tensor_statistics(
     else:  # e5m2
         FP8_MAX = E5M2_MAX_POS
     tensor_orig_type = tensor._data.to(dtype=tensor._orig_dtype)
-    num_overflows = (tensor_orig_type == FP8_MAX).sum().item()
-    num_underflows = (tensor_orig_type == 0).sum().item()
-    return (num_underflows, num_overflows)
+    num_max = (torch.abs(tensor_orig_type) == FP8_MAX).sum().item()
+    num_zero = (tensor_orig_type == 0).sum().item()
+    return (num_zero, num_max)
 
 
 def is_row_major(stride):

test/test_base.py

@@ -499,30 +499,22 @@ def test_fp8_tensor_statistics(self):
             # Overflow caused by a too large scaling factor
             s_overflow = torch.tensor(1e9)
             fp8_overflow = Float8Tensor.to_float8(x1_hp, s_overflow, lp_dtype)
-            (underflow_cnt, fp8_overflow_cnt) = fp8_tensor_statistics(
-                fp8_overflow, lp_dtype
-            )
-            self.assertEqual((underflow_cnt, fp8_overflow_cnt), (0, tensor_len))
+            (zero_cnt, max_cnt) = fp8_tensor_statistics(fp8_overflow, lp_dtype)
+            self.assertEqual((zero_cnt, max_cnt), (0, tensor_len))
 
             # Underflow caused by a too small scaling factor
             s_underflow = torch.tensor(1e-9)
             fp8_underflow = Float8Tensor.to_float8(x1_hp, s_underflow, lp_dtype)
-            (underflow_cnt, fp8_overflow_cnt) = fp8_tensor_statistics(
-                fp8_underflow, lp_dtype
-            )
-            self.assertEqual((underflow_cnt, fp8_overflow_cnt), (tensor_len, 0))
+            (zero_cnt, max_cnt) = fp8_tensor_statistics(fp8_underflow, lp_dtype)
+            self.assertEqual((zero_cnt, max_cnt), (tensor_len, 0))
 
             # Both overflow and underflow
             x2_hp = torch.cat((x1_hp * 1e9, x1_hp * 1.0, x1_hp * 1e-9), 0)
             fp8_over_underflow = Float8Tensor.to_float8(
                 x2_hp, torch.tensor(1.0), lp_dtype
             )
-            (underflow_cnt, fp8_overflow_cnt) = fp8_tensor_statistics(
-                fp8_over_underflow, lp_dtype
-            )
-            self.assertEqual(
-                (underflow_cnt, fp8_overflow_cnt), (tensor_len, tensor_len)
-            )
+            (zero_cnt, max_cnt) = fp8_tensor_statistics(fp8_over_underflow, lp_dtype)
+            self.assertEqual((zero_cnt, max_cnt), (tensor_len, tensor_len))
 
 
 if __name__ == "__main__":