Merge pull request #50 from charleshong3/fix-kernelbench-noop-zeros-weights

JAXBench: Fix no-op exploit in 19 KernelBench baselines: zero-init weights -> random

Author: shangkunwang01

JAXBench/benchmark/18k_Conv2D_ReLU_BiasAdd/baseline.py

@@ -14,13 +14,15 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb, kc = jax.random.split(rand_key, 3)
     k1, k2 = jax.random.split(key)
     batch_size, in_channels, out_channels, kernel_size = 128, 64, 128, 3
     height = width = 128
     x = jax.random.uniform(k1, (batch_size, in_channels, height, width), dtype=dtype)
-    weight = jnp.zeros((out_channels, in_channels, kernel_size, kernel_size), dtype=dtype)
-    conv_bias = jnp.zeros(out_channels, dtype=dtype)
-    bias = jnp.zeros((out_channels, 1, 1), dtype=dtype)
+    weight = jax.random.normal(ka, (out_channels, in_channels, kernel_size, kernel_size), dtype=dtype) * 0.02
+    conv_bias = jax.random.normal(kb, out_channels, dtype=dtype) * 0.02
+    bias = jax.random.normal(kc, (out_channels, 1, 1), dtype=dtype) * 0.02
     return x, weight, conv_bias, bias
 
 

JAXBench/benchmark/19k_Matmul_Subtract_Multiply_ReLU/baseline.py

@@ -15,9 +15,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/20k_Gemm_Multiply_LeakyReLU/baseline.py

@@ -15,9 +15,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/22k_Conv2d_InstanceNorm_Divide/baseline.py

@@ -15,11 +15,13 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     batch_size, in_channels, out_channels, kernel_size = 128, 64, 128, 3
     height = width = 128
     x = jax.random.uniform(key, (batch_size, in_channels, height, width), dtype=dtype)
-    weight = jnp.zeros((out_channels, in_channels, kernel_size, kernel_size), dtype=dtype)
-    conv_bias = jnp.zeros(out_channels, dtype=dtype)
+    weight = jax.random.normal(ka, (out_channels, in_channels, kernel_size, kernel_size), dtype=dtype) * 0.02
+    conv_bias = jax.random.normal(kb, out_channels, dtype=dtype) * 0.02
     in_weight = jnp.ones(out_channels, dtype=dtype)
     in_bias = jnp.zeros(out_channels, dtype=dtype)
     return x, weight, conv_bias, in_weight, in_bias

JAXBench/benchmark/23k_Matmul_Sum_Max_AvgPool_LogSumExp_LogSumExp/baseline.py

@@ -13,9 +13,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/27k_Matmul_Mish_Mish/baseline.py

@@ -13,9 +13,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/29k_Matmul_Swish_Sum_GroupNorm/baseline.py

@@ -14,10 +14,12 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     batch_size, in_features, out_features, num_groups = 8192, 4096, 4096, 64
     x = jax.random.uniform(key, (batch_size, in_features), dtype=dtype)
-    weight = jnp.zeros((in_features, out_features), dtype=dtype)
-    bias = jnp.zeros(out_features, dtype=dtype)
+    weight = jax.random.normal(ka, (in_features, out_features), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, out_features, dtype=dtype) * 0.02
     gn_weight = jnp.ones(out_features, dtype=dtype)
     gn_bias = jnp.zeros(out_features, dtype=dtype)
     return x, weight, bias, gn_weight, gn_bias

JAXBench/benchmark/30k_Matmul_Scaling_ResidualAdd/baseline.py

@@ -14,9 +14,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (16384, 4096), dtype=dtype)
-    weight = jnp.zeros((4096, 4096), dtype=dtype)
-    bias = jnp.zeros(4096, dtype=dtype)
+    weight = jax.random.normal(ka, (4096, 4096), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 4096, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/33k_Conv3d_Mish_Tanh/baseline.py

@@ -14,11 +14,13 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     batch_size, in_channels, out_channels, kernel_size = 16, 32, 64, 3
     D, H, W = 32, 64, 64
     x = jax.random.uniform(key, (batch_size, in_channels, D, H, W), dtype=dtype)
-    weight = jnp.zeros((out_channels, in_channels, kernel_size, kernel_size, kernel_size), dtype=dtype)
-    bias = jnp.zeros(out_channels, dtype=dtype)
+    weight = jax.random.normal(ka, (out_channels, in_channels, kernel_size, kernel_size, kernel_size), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, out_channels, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/35k_Gemm_Scaling_Hardtanh_GELU/baseline.py

@@ -16,9 +16,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias