fix

Signed-off-by: Pawel Gadzinski <[email protected]>

Author: pggPL

tests/pytorch/debug/test_log.py

@@ -164,8 +164,7 @@ def test_numerics(fp8_recipe, feature_dirs):
             num_quantizers=3,
         )
 
-        tensor = torch.zeros(1024, 1024).cuda()
-        tensor[0, :] = 1000
+        tensor = torch.randn(1024, 1024).cuda()
         quantizer = recipe_state.make_quantizers()[0]
         quantized_tensor = quantizer(tensor)
 
@@ -189,14 +188,14 @@ def test_numerics(fp8_recipe, feature_dirs):
             underflows = float(line.split("value=")[1])
             expected = (
                 ((dequantized_tensor == 0).sum() - (tensor == 0).sum())
-                / dequantized_tensor.numel()
+                / tensor.numel()
                 * 100
             )
             assert underflows == pytest.approx(expected.cpu(), abs=1e-4)
         if "mse" in line:
             mse = float(line.split("value=")[1])
             expected = torch.nn.functional.mse_loss(dequantized_tensor, tensor, reduction="mean")
-            assert mse == pytest.approx(expected.cpu(), abs=1e-6)
+            assert mse == pytest.approx(expected.cpu(), abs=1e-4)
         if "overflows%" in line:
             overflows = float(line.split("value=")[1])
             expected = (

transformer_engine/debug/features/utils/stats_computation.py

@@ -211,20 +211,20 @@ def add_underflows_stats(recipe_name: str, columnwise: bool = False):
     stats_to_num[stat_num] = len(stats_to_num)
     stats_to_num[stat_pct] = len(stats_to_num)
 
-    zero_values = torch.tensor([0, 127], device="cuda")
+    zero_values = torch.tensor([0, 128], device="cuda")
 
     STATS[stat_num] = (
         lambda x, aux_dict: 
-            aux_dict[recipe_name].get_data_tensors(
+            torch.isin(aux_dict[recipe_name].get_data_tensors(
                 rowwise_data=not columnwise, columnwise_data=columnwise
-            ).isin(zero_values).sum() - (x == 0).sum(),
+            ), zero_values).sum() - (x == 0).sum(),
         lambda buffers, _sn=stat_num: sum(_get(buffers, _sn)),
     )
     STATS[stat_pct] = (
         lambda x, aux_dict: (
-            aux_dict[recipe_name].get_data_tensors(
+            torch.isin(aux_dict[recipe_name].get_data_tensors(
                 rowwise_data=not columnwise, columnwise_data=columnwise
-            ).isin(zero_values).sum() - (x == 0).sum())
+            ), zero_values).sum() - (x == 0).sum())
         / aux_dict[recipe_name].numel()
         * 100,
         lambda buffers, _sn_num=stat_num: 100