Update int 4 flow for consistency

Summary: Update int 4 flow for consistency

Reviewed By: kimishpatel

Differential Revision: D54350318

fbshipit-source-id: 7d04c46f2fc78e277b515828f2a47f4acb5c7e86

Author: Michael Gschwind

examples/models/llama2/export_llama_lib.py

@@ -192,11 +192,8 @@ def quantize(
     elif qmode == "int4":
         model_int4 = Int8DynActInt4WeightQuantHandler(
             model, activation_precision=torch_dtype
-        )
-        model_int4_state_dict = model_int4.create_quantized_state_dict()
-        model_int4 = model_int4.convert_for_runtime()
+        ).quantized_model()
         print("quantized model:", model_int4)
-        model_int4.load_state_dict(model_int4_state_dict)
         return model_int4
     else:
         raise Exception(f"Unrecognized quantize mode: {qmode}")

examples/models/llama2/quantize.py

@@ -1002,6 +1002,12 @@ def convert_for_runtime(self):
         )
         return self.mod
 
+    def quantized_model(self) -> nn.Module:
+        model_updated_state_dict = self.create_quantized_state_dict()
+        self.convert_for_runtime()
+        self.mod.load_state_dict(model_updated_state_dict)
+        return self.mod
+
 
 class Int8DynActInt4WeightLinear(torch.nn.Module):
     __constants__ = ["in_features", "out_features"]