Updates

Author: jainapurva

benchmarks/microbenchmarks/test/benchmark_config.yml

@@ -8,8 +8,8 @@ quantization_config_recipe_names:
   - "float8wo"
 # sparsity_config_recipe_names:
   # Will run a baseline inference for model by default, without sparsity for comparison
-  - "semi-sparse"
-  - "block"
+  # - "semi-sparse"
+  # - "block"
 output_dir: "benchmarks/microbenchmarks/results"
 model_params:
   - name: "small_bf16_linear"

test/test_model_architecture.py

@@ -6,22 +6,47 @@
 
 import unittest
 
+import torch
+from parameterized import parameterized
+
 from torchao.testing.model_architectures import create_model_and_input_data
+from torchao.utils import get_available_devices
 
 
 class TestModels(unittest.TestCase):
-    def test_toy_linear_model(self):
-        model, input_data = create_model_and_input_data("linear", 10, 64, 32)
+    @parameterized.expand([(device,) for device in get_available_devices()])
+    def test_toy_linear_model(self, device):
+        # Skip if device is not available
+        if device == "cuda" and not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+
+        model, input_data = create_model_and_input_data(
+            "linear", 10, 64, 32, device=device
+        )
         output = model(input_data)
         self.assertEqual(output.shape, (10, 32))
 
-    def test_ln_linear_activation_model(self):
-        model, input_data = create_model_and_input_data("ln_linear_sigmoid", 10, 64, 32)
+    @parameterized.expand([(device,) for device in get_available_devices()])
+    def test_ln_linear_activation_model(self, device):
+        # Skip if device is not available
+        if device == "cuda" and not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+
+        model, input_data = create_model_and_input_data(
+            "ln_linear_sigmoid", 10, 64, 32, device=device
+        )
         output = model(input_data)
         self.assertEqual(output.shape, (10, 32))
 
-    def test_transformer_block(self):
-        model, input_data = create_model_and_input_data("transformer_block", 10, 64, 32)
+    @parameterized.expand([(device,) for device in get_available_devices()])
+    def test_transformer_block(self, device):
+        # Skip if device is not available
+        if device == "cuda" and not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+
+        model, input_data = create_model_and_input_data(
+            "transformer_block", 10, 64, 32, device=device
+        )
         output = model(input_data)
         self.assertEqual(output.shape, (10, 16, 64))
 

torchao/testing/model_architectures.py

@@ -21,9 +21,7 @@ def forward(self, x):
 
 
 class LNLinearActivationModel(nn.Module):
-    def __init__(
-        self, fc_dim1, fc_dim2, dtype=torch.bfloat16, activation="sigmoid", device=None
-    ):
+    def __init__(self, fc_dim1, fc_dim2, dtype=torch.bfloat16, activation="sigmoid"):
         super().__init__()
 
         activation = activation.lower()
@@ -41,7 +39,7 @@ def __init__(
             raise ValueError(f"Unsupported activation: {activation}")
 
         self.ln = nn.LayerNorm(fc_dim1, elementwise_affine=False)
-        self.fc = nn.Linear(fc_dim1, fc_dim2, bias=False).to(dtype=dtype, device=device)
+        self.fc = nn.Linear(fc_dim1, fc_dim2, bias=False).to(dtype=dtype)
         self.activation = activation_map[activation]
 
     def forward(self, x):
@@ -50,6 +48,20 @@ def forward(self, x):
         return self.activation(x)
 
 
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
 class TransformerBlock(torch.nn.Module):
     def __init__(self, hidden_dim, num_heads=8, mlp_ratio=4, dtype=torch.bfloat16):
         super().__init__()
@@ -72,8 +84,8 @@ def __init__(self, hidden_dim, num_heads=8, mlp_ratio=4, dtype=torch.bfloat16):
         )
 
         # Layer norms
-        self.norm1 = nn.RMSNorm(hidden_dim, dtype=dtype)
-        self.norm2 = nn.RMSNorm(hidden_dim, dtype=dtype)
+        self.norm1 = RMSNorm(hidden_dim).to(dtype)
+        self.norm2 = RMSNorm(hidden_dim).to(dtype)
 
         # Activation
         self.activation = torch.nn.GELU()