Replace torch.empty with torch.zeros (pytorch#6478)

Summary:

X-link: pytorch/ao#1157

It turns out that it is unsafe to use `torch.empty` in oss environment because `torch.empty` creates tensor with uninitialized data. That means the buffer could be initialized with random values depends on what is left on that piece of memory. This causes code to generate inconsistent behavior.

This PR replaces `torch.empty` with `torch.zeros` to make sure that they are properly initialized and avoid inconsistent behaviors.

Reviewed By: msaroufim

Differential Revision: D64875312

Author: Lunwen He

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
+        super().zeros = torch.zeros_like(super().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,