Update

[ghstack-poisoned]

Author: jainapurva

benchmarks/microbenchmarks/test/test_benchmark_runner.py

@@ -63,40 +63,56 @@ def test_get_shapes_for_config(self):
         )
         self.assertEqual(len(shapes), 1)
         self.assertEqual(shapes[0], ("custom", [1024, 1024, 1024]))
-        
+
         # Test llama shapes
-        llama_shapes = get_shapes_for_config([
-            {"name": "llama"}
-        ])
+        llama_shapes = get_shapes_for_config([{"name": "llama"}])
         self.assertEqual(len(llama_shapes), 4)  # 4 LLaMa shapes
-        self.assertTrue(any(name.startswith("llama_attn.wqkv") for name, _ in llama_shapes))
-        self.assertTrue(any(name.startswith("llama_attn.w0") for name, _ in llama_shapes))
-        self.assertTrue(any(name.startswith("llama_ffn.w13") for name, _ in llama_shapes))
-        self.assertTrue(any(name.startswith("llama_ffn.w2") for name, _ in llama_shapes))
-        
+        self.assertTrue(
+            any(name.startswith("llama_attn.wqkv") for name, _ in llama_shapes)
+        )
+        self.assertTrue(
+            any(name.startswith("llama_attn.w0") for name, _ in llama_shapes)
+        )
+        self.assertTrue(
+            any(name.startswith("llama_ffn.w13") for name, _ in llama_shapes)
+        )
+        self.assertTrue(
+            any(name.startswith("llama_ffn.w2") for name, _ in llama_shapes)
+        )
+
         # Test pow2 shapes
-        pow2_shapes = get_shapes_for_config([
-            {"name": "pow2", "min_power": 10, "max_power": 12}
-        ])
+        pow2_shapes = get_shapes_for_config(
+            [{"name": "pow2", "min_power": 10, "max_power": 12}]
+        )
         self.assertEqual(len(pow2_shapes), 3)  # 3 powers of 2 (10, 11, 12)
         self.assertEqual(pow2_shapes[0], ("pow2_0", [1024, 1024, 1024]))  # 2^10
         self.assertEqual(pow2_shapes[1], ("pow2_1", [2048, 2048, 2048]))  # 2^11
         self.assertEqual(pow2_shapes[2], ("pow2_2", [4096, 4096, 4096]))  # 2^12
-        
+
         # Test pow2_extended shapes
-        pow2_extended_shapes = get_shapes_for_config([
-            {"name": "pow2_extended", "min_power": 10, "max_power": 11}
-        ])
-        self.assertEqual(len(pow2_extended_shapes), 4)  # 2 powers of 2, each with 2 variants
-        self.assertEqual(pow2_extended_shapes[0], ("pow2_extended_0", [1024, 1024, 1024]))  # 2^10
-        self.assertEqual(pow2_extended_shapes[1], ("pow2_extended_1", [1536, 1536, 1536]))  # 2^10 + 2^9
-        self.assertEqual(pow2_extended_shapes[2], ("pow2_extended_2", [2048, 2048, 2048]))  # 2^11
-        self.assertEqual(pow2_extended_shapes[3], ("pow2_extended_3", [3072, 3072, 3072]))  # 2^11 + 2^10
-        
+        pow2_extended_shapes = get_shapes_for_config(
+            [{"name": "pow2_extended", "min_power": 10, "max_power": 11}]
+        )
+        self.assertEqual(
+            len(pow2_extended_shapes), 4
+        )  # 2 powers of 2, each with 2 variants
+        self.assertEqual(
+            pow2_extended_shapes[0], ("pow2_extended_0", [1024, 1024, 1024])
+        )  # 2^10
+        self.assertEqual(
+            pow2_extended_shapes[1], ("pow2_extended_1", [1536, 1536, 1536])
+        )  # 2^10 + 2^9
+        self.assertEqual(
+            pow2_extended_shapes[2], ("pow2_extended_2", [2048, 2048, 2048])
+        )  # 2^11
+        self.assertEqual(
+            pow2_extended_shapes[3], ("pow2_extended_3", [3072, 3072, 3072])
+        )  # 2^11 + 2^10
+
         # Test sweep shapes (limited to a small range for testing)
-        sweep_shapes = get_shapes_for_config([
-            {"name": "sweep", "min_power": 8, "max_power": 9}
-        ])
+        sweep_shapes = get_shapes_for_config(
+            [{"name": "sweep", "min_power": 8, "max_power": 9}]
+        )
         # For min_power=8, max_power=9, we should have 8 shapes (2^3 = 8 combinations)
         self.assertEqual(len(sweep_shapes), 8)
         # Check that all shapes have the expected format

benchmarks/microbenchmarks/test/test_utils.py

@@ -171,7 +171,7 @@ def test_rms_norm(self):
         x = torch.randn(16, 64)
         out = rms_norm(x)
         self.assertEqual(out.shape, (16, 64))
-        
+
         # Test with different eps
         rms_norm = RMSNorm(dim=64, eps=1e-5)
         out = rms_norm(x)
@@ -184,38 +184,50 @@ def test_rms_norm_linear_activation(self):
         out = model(x)
         self.assertEqual(out.shape, (16, 32))
         self.assertEqual(out.dtype, torch.float32)
-        
+
         # Test with ReLU activation
-        model = RMSNormLinearActivation(fc_dim1=64, fc_dim2=32, dtype=torch.float32, activation="relu")
+        model = RMSNormLinearActivation(
+            fc_dim1=64, fc_dim2=32, dtype=torch.float32, activation="relu"
+        )
         out = model(x)
         self.assertEqual(out.shape, (16, 32))
         self.assertTrue(torch.all(out >= 0))  # Check ReLU output range
-        
+
         # Test with SiLU activation
-        model = RMSNormLinearActivation(fc_dim1=64, fc_dim2=32, dtype=torch.float32, activation="silu")
+        model = RMSNormLinearActivation(
+            fc_dim1=64, fc_dim2=32, dtype=torch.float32, activation="silu"
+        )
         out = model(x)
         self.assertEqual(out.shape, (16, 32))
-        
+
         # Test with invalid activation
         with self.assertRaises(ValueError):
-            RMSNormLinearActivation(fc_dim1=64, fc_dim2=32, dtype=torch.float32, activation="invalid")
+            RMSNormLinearActivation(
+                fc_dim1=64, fc_dim2=32, dtype=torch.float32, activation="invalid"
+            )
 
     def test_transformer_block(self):
         # Test with default parameters
-        model = TransformerBlock(hidden_dim=64, num_heads=8, mlp_ratio=4, dtype=torch.float32)
+        model = TransformerBlock(
+            hidden_dim=64, num_heads=8, mlp_ratio=4, dtype=torch.float32
+        )
         x = torch.randn(16, 16, 64)  # [batch_size, seq_len, hidden_dim]
         out = model(x)
         self.assertEqual(out.shape, (16, 16, 64))
         self.assertEqual(out.dtype, torch.float32)
-        
+
         # Test with different parameters
-        model = TransformerBlock(hidden_dim=128, num_heads=4, mlp_ratio=2, dtype=torch.float32)
+        model = TransformerBlock(
+            hidden_dim=128, num_heads=4, mlp_ratio=2, dtype=torch.float32
+        )
         x = torch.randn(8, 32, 128)
         out = model(x)
         self.assertEqual(out.shape, (8, 32, 128))
-        
+
         # Test with different head dimensions
-        model = TransformerBlock(hidden_dim=96, num_heads=6, mlp_ratio=3, dtype=torch.float32)
+        model = TransformerBlock(
+            hidden_dim=96, num_heads=6, mlp_ratio=3, dtype=torch.float32
+        )
         x = torch.randn(4, 8, 96)
         out = model(x)
         self.assertEqual(out.shape, (4, 8, 96))
@@ -255,7 +267,7 @@ def test_create_model_and_input(self):
         )
         self.assertIsInstance(model, RMSNormLinearActivation)
         self.assertEqual(input_data.shape, (m, k))
-        
+
         # Test TransformerBlock
         model, input_data = create_model_and_input(
             model_type="transformer_block",
@@ -266,40 +278,50 @@ def test_create_model_and_input(self):
             device="cpu",
         )
         self.assertIsInstance(model, TransformerBlock)
-        self.assertEqual(input_data.shape, (m, 16, k))  # [batch_size, seq_len, hidden_dim]
+        self.assertEqual(
+            input_data.shape, (m, 16, k)
+        )  # [batch_size, seq_len, hidden_dim]
 
     def test_quantization_on_models(self):
         # Test quantization on RMSNormLinearActivation
         model = RMSNormLinearActivation(fc_dim1=64, fc_dim2=32, dtype=torch.float32)
         x = torch.randn(16, 64)
-        
+
         # Test with Int8WeightOnlyConfig
         config = string_to_config(quantization="int8wo", sparsity=None)
         if config is not None:
             # Skip quantization test if torchao.quantization.quantize is not available
             try:
                 from torchao.quantization import quantize
+
                 quantized_model = quantize(model, config)
                 out = quantized_model(x)
                 self.assertEqual(out.shape, (16, 32))
             except ImportError:
-                print("Skipping quantization test: torchao.quantization.quantize not available")
-        
+                print(
+                    "Skipping quantization test: torchao.quantization.quantize not available"
+                )
+
         # Test quantization on TransformerBlock
-        model = TransformerBlock(hidden_dim=64, num_heads=8, mlp_ratio=4, dtype=torch.float32)
+        model = TransformerBlock(
+            hidden_dim=64, num_heads=8, mlp_ratio=4, dtype=torch.float32
+        )
         x = torch.randn(16, 16, 64)
-        
+
         # Test with Int8WeightOnlyConfig
         config = string_to_config(quantization="int8wo", sparsity=None)
         if config is not None:
             # Skip quantization test if torchao.quantization.quantize is not available
             try:
                 from torchao.quantization import quantize
+
                 quantized_model = quantize(model, config)
                 out = quantized_model(x)
                 self.assertEqual(out.shape, (16, 16, 64))
             except ImportError:
-                print("Skipping quantization test: torchao.quantization.quantize not available")
+                print(
+                    "Skipping quantization test: torchao.quantization.quantize not available"
+                )
 
     def test_generate_results_csv(self):
         results = [

benchmarks/microbenchmarks/utils.py

@@ -399,7 +399,7 @@ def __init__(self, fc_dim1, fc_dim2, dtype=torch.bfloat16, activation="gelu"):
         super().__init__()
         self.rms_norm = RMSNorm(fc_dim1, dtype=dtype)
         self.fc = torch.nn.Linear(fc_dim1, fc_dim2, bias=False).to(dtype)
-        
+
         if activation == "gelu":
             self.activation = torch.nn.GELU()
         elif activation == "relu":
@@ -422,66 +422,72 @@ def __init__(self, hidden_dim, num_heads=8, mlp_ratio=4, dtype=torch.bfloat16):
         self.hidden_dim = hidden_dim
         self.num_heads = num_heads
         self.head_dim = hidden_dim // num_heads
-        
+
         # Self-attention
         self.qkv = torch.nn.Linear(hidden_dim, 3 * hidden_dim, bias=False).to(dtype)
         self.proj = torch.nn.Linear(hidden_dim, hidden_dim, bias=False).to(dtype)
-        
+
         # MLP
         self.mlp_ratio = mlp_ratio
         self.mlp_hidden_dim = int(hidden_dim * mlp_ratio)
-        self.mlp_fc1 = torch.nn.Linear(hidden_dim, self.mlp_hidden_dim, bias=False).to(dtype)
-        self.mlp_fc2 = torch.nn.Linear(self.mlp_hidden_dim, hidden_dim, bias=False).to(dtype)
-        
+        self.mlp_fc1 = torch.nn.Linear(hidden_dim, self.mlp_hidden_dim, bias=False).to(
+            dtype
+        )
+        self.mlp_fc2 = torch.nn.Linear(self.mlp_hidden_dim, hidden_dim, bias=False).to(
+            dtype
+        )
+
         # Layer norms
         self.norm1 = RMSNorm(hidden_dim, dtype=dtype)
         self.norm2 = RMSNorm(hidden_dim, dtype=dtype)
-        
+
         # Activation
         self.activation = torch.nn.GELU()
 
     def forward(self, x):
         batch_size, seq_len, _ = x.shape
-        
+
         # Self-attention
         residual = x
         x = self.norm1(x)
-        
+
         # Reshape qkv projection for better memory layout
         qkv = self.qkv(x)  # [batch_size, seq_len, 3 * hidden_dim]
         qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
-        qkv = qkv.permute(2, 0, 3, 1, 4)  # [3, batch_size, num_heads, seq_len, head_dim]
+        qkv = qkv.permute(
+            2, 0, 3, 1, 4
+        )  # [3, batch_size, num_heads, seq_len, head_dim]
         q, k, v = qkv  # Each has shape [batch_size, num_heads, seq_len, head_dim]
-        
+
         # Scaled dot-product attention with proper reshaping
         # Reshape for better memory layout and avoid broadcasting issues
         q = q.reshape(batch_size * self.num_heads, seq_len, self.head_dim)
         k = k.reshape(batch_size * self.num_heads, seq_len, self.head_dim)
         v = v.reshape(batch_size * self.num_heads, seq_len, self.head_dim)
-        
+
         # Compute attention scores
-        attn = (q @ k.transpose(-2, -1)) * (1.0 / (self.head_dim ** 0.5))
+        attn = (q @ k.transpose(-2, -1)) * (1.0 / (self.head_dim**0.5))
         attn = torch.softmax(attn, dim=-1)
-        
+
         # Apply attention to values
         x = attn @ v  # [batch_size * num_heads, seq_len, head_dim]
-        
+
         # Reshape back to original dimensions
         x = x.reshape(batch_size, self.num_heads, seq_len, self.head_dim)
         x = x.transpose(1, 2).reshape(batch_size, seq_len, self.hidden_dim)
-        
+
         # Project back to hidden dimension
         x = self.proj(x)
         x = residual + x
-        
+
         # MLP
         residual = x
         x = self.norm2(x)
         x = self.mlp_fc1(x)
         x = self.activation(x)
         x = self.mlp_fc2(x)
         x = residual + x
-        
+
         return x
 
 
@@ -683,7 +689,9 @@ def create_model_and_input(
         input_data = torch.randn(m, k, device=device, dtype=high_precision_dtype)
     elif model_type == "transformer_block":
         # For transformer block, k is the hidden dimension
-        model = TransformerBlock(k, num_heads=8, mlp_ratio=4, dtype=high_precision_dtype).to(device)
+        model = TransformerBlock(
+            k, num_heads=8, mlp_ratio=4, dtype=high_precision_dtype
+        ).to(device)
         # Input shape for transformer is [batch_size, seq_len, hidden_dim]
         input_data = torch.randn(m, 16, k, device=device, dtype=high_precision_dtype)
     else: