Replace torch.empty with torch.zeros

Differential Revision: D64875312

Pull Request resolved: #6478

Author: helunwencser

examples/models/llama/source_transformation/lora.py

@@ -70,6 +70,8 @@ def __init__(
             precision=precision,
             scales_precision=scales_precision,
         )
+        # TODO(lunwenh): Remove this once TorchAO's commit pin in ExecuTorch is updated to include this PR
+        self.zeros = torch.zeros_like(self.zeros)
         self.adaptor = LoRAAdaptorLinear(
             in_features,
             out_features,

examples/models/llama/source_transformation/pre_quantization.py

@@ -46,6 +46,8 @@ def replacement_fn(child: torch.nn.Module) -> torch.nn.Module:
             precision=precision,
             scales_precision=scales_precision,
         )
+        # TODO(lunwenh): Remove this once TorchAO's commit pin in ExecuTorch is updated to include this PR
+        new_linear.zeros = torch.zeros_like(new_linear.zeros)
         return new_linear
 
     _replace_with_custom_fn_if_matches_filter(module, replacement_fn, filter_fn)

examples/models/llama/source_transformation/quantize.py

@@ -375,7 +375,7 @@ def __init__(
         self.in_features = in_features
         self.out_features = out_features
         self.register_buffer(
-            "weight", torch.empty((out_features, in_features), dtype=torch.int8)
+            "weight", torch.zeros((out_features, in_features), dtype=torch.int8)
         )
         self.register_buffer("scales", torch.ones(out_features, dtype=torch.bfloat16))
 
@@ -448,18 +448,18 @@ def __init__(
         # currently storing unpacked int8 weights
         self.register_buffer(
             "weight",
-            torch.empty((out_features, in_features), dtype=torch.int8),
+            torch.zeros((out_features, in_features), dtype=torch.int8),
         )
         self.register_buffer(
             "scales",
-            torch.empty(
+            torch.zeros(
                 (out_features),
                 dtype=torch.float32,
             ),
         )
         self.register_buffer(
             "zeros",
-            torch.empty(
+            torch.zeros(
                 (out_features),
                 dtype=torch.float32,
             ),
@@ -632,15 +632,15 @@ def __init__(
         if not packed:
             self.register_buffer(
                 "weight",
-                torch.empty(
+                torch.zeros(
                     (vocab_size, embedding_dim), dtype=torch.int8, device=device
                 ),
             )
         else:  # packed
             if bitwidth == 2:
                 self.register_buffer(
                     "weight",
-                    torch.empty(
+                    torch.zeros(
                         (vocab_size, embedding_dim // 4),
                         dtype=torch.uint8,
                         device=device,
@@ -649,7 +649,7 @@ def __init__(
             elif bitwidth == 4:
                 self.register_buffer(
                     "weight",
-                    torch.empty(
+                    torch.zeros(
                         (vocab_size, embedding_dim // 2),
                         dtype=torch.uint8,
                         device=device,