Update on "redo of add quantized layer norm implementation"

vkuzo · vkuzo · commit 18eb6cdf87db · 2020-04-15T08:39:46.000-07:00
Summary: This is a redo of #35329 with a better test. Adds a quantized implementation of LayerNorm for server. A future PR will add the Python wrapper. Test Plan: numerics match the floating point implementation benchmarks by input size: v1 (mean+var non-vectorized): https://gist.github.com/vkuzo/f6d72c04742608112f4c2e612c74bd13 v2 (mean+var vectorized in float): https://gist.github.com/vkuzo/4dd95657c5b5f3654e0965db00eff8d2 v3 (mean+var vectorized in int, current): https://gist.github.com/vkuzo/57a75f75629da9f23b64b38ca0e3d34b Differential Revision: [D21030268](https://our.internmc.facebook.com/intern/diff/D21030268) [ghstack-poisoned]
diff --git a/test/quantization/test_quantized.py b/test/quantization/test_quantized.py
@@ -356,15 +356,11 @@ def test_qlayer_norm(self, shapes, torch_type, X_rand_scale, Y_scale, Y_zero_poi
             dqX = qX.dequantize()
 
             # Enforce non-homogeneous inputs
-            nonzero_var_in_each_layer = sum(
-                1 if ((dqX[i] - dqX[i].min()) / (dqX[i].max() - dqX[i].min() + 1e-5)).std() > 1e-2 else 0
-                for i in range(dqX.shape[0])
-            ) == dqX.shape[0]
-            assume(nonzero_var_in_each_layer)
             enough_unique_vals_in_each_layer = sum(
                 1 if (
                     dqX[i].shape[0] < 5 or
-                    float(torch.unique(dqX[i]).shape[0]) / dqX[i].shape[0] > 0.01) else 0
+                    float(torch.unique(dqX[i]).shape[0]) / dqX[i].shape[0] > 0.01
+                ) else 0
                 for i in range(dqX.shape[0])
             ) == dqX.shape[0]
             assume(enough_unique_vals_in_each_layer)
