Remove exir.capture and its warnings

Summary: Remove exir.capture and its warnings

Differential Revision: D55039873

Author: Michael Gschwind

examples/models/llama2/quantize.py

@@ -247,7 +247,7 @@ class WeightOnlyInt8Linear(torch.nn.Module):
     __constants__ = ["in_features", "out_features"]
     in_features: int
     out_features: int
-    weight: torch.Tensor
+    # weight: torch.Tensor
 
     def __init__(
         self,
@@ -260,10 +260,15 @@ def __init__(
         super().__init__()
         self.in_features = in_features
         self.out_features = out_features
-        self.register_buffer(
-            "weight", torch.empty((out_features, in_features), dtype=torch.int8)
+        self.register_parameter(
+            "weight",
+            torch.nn.Parameter(
+                torch.empty((out_features, in_features), dtype=torch.int8)
+            ),
+        )
+        self.register_parameter(
+            "scales", torch.nn.Parameter(torch.ones(out_features, dtype=torch.bfloat16))
         )
-        self.register_buffer("scales", torch.ones(out_features, dtype=torch.bfloat16))
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         return F.linear(input, self.weight.to(dtype=input.dtype)) * self.scales
@@ -372,17 +377,24 @@ def __init__(
             group_size = embedding_dim
         self.group_size = group_size
         self.dtype = dtype
-        self.register_buffer(
-            "weight", torch.empty((vocab_size, embedding_dim), dtype=torch.int8)
+        self.register_parameter(
+            "weight",
+            torch.nn.Parameter(
+                torch.empty((vocab_size, embedding_dim), dtype=torch.int8)
+            ),
         )
         groups_per_row = (embedding_dim + group_size - 1) // group_size
         if groups_per_row > 1:
-            self.register_buffer(
-                "scales", torch.ones((vocab_size, groups_per_row), dtype=torch.float16)
+            self.register_parameter(
+                "scales",
+                torch.nn.Parameter(
+                    torch.ones((vocab_size, groups_per_row), dtype=torch.float16)
+                ),
             )
         else:
-            self.register_buffer(
-                "scales", torch.ones((vocab_size,), dtype=torch.float16)
+            self.register_parameter(
+                "scales",
+                torch.nn.Parameter(torch.ones((vocab_size,), dtype=torch.float16)),
             )
 
     @torch.no_grad()
@@ -583,7 +595,7 @@ class Int8DynActInt4WeightLinear(torch.nn.Module):
 
     in_features: int
     out_features: int
-    weight: torch.Tensor
+    # weight: torch.Tensor
 
     """
     This module implements a dynamic quantized linear layer with int4 weight.
@@ -624,28 +636,30 @@ def __init__(
         self.precision = precision
 
         # currently storing unpacked int8 weights
-        # TODO: ????!!!!!
-        # weights should be registers as parameters, since they're
-        # read-only for inference
-        self.register_buffer(
+        self.register_parameter(
             "weight",
-            torch.empty((out_features, in_features), dtype=torch.int8),
+            torch.nn.Parameter(
+                torch.empty((out_features, in_features), dtype=torch.int8)
+            ),
         )
-        self.register_buffer(
+        self.register_parameter(
             "scales",
-            torch.empty(
-                (out_features, in_features // group_size),
-                dtype=scales_precision,
+            torch.nn.Parameter(
+                torch.empty(
+                    (out_features, in_features // group_size), dtype=scales_precision
+                ),
             ),
         )
         # TODO:
         # Let's not store 0 - and then have to process them?!
         # All our quantization is symmetric.
-        self.register_buffer(
+        self.register_parameter(
             "zeros",
-            torch.empty(
-                (out_features, in_features // group_size),
-                dtype=scales_precision,
+            torch.nn.Parameter(
+                torch.empty(
+                    (out_features, in_features // group_size),
+                    dtype=scales_precision,
+                )
             ),
         )