up

Author: metascroy

examples/models/llama/CMakeLists.txt

@@ -116,16 +116,20 @@ endif()
 
 if(EXECUTORCH_BUILD_TORCHAO)
   set(TORCHAO_BUILD_EXECUTORCH_OPS ON)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../../../third-party/ao/torchao/experimental ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/ao/torchao/experimental)
+  set(TORCHAO_BUILD_CPU_AARCH64 ON)
+  add_subdirectory(
+    ${CMAKE_CURRENT_SOURCE_DIR}/../../../third-party/ao/torchao/experimental
+    ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/ao/torchao/experimental
+  )
   target_link_options_shared_lib(torchao_ops_executorch)
   list(APPEND link_libraries torchao_ops_executorch)
-  if(CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
-    add_subdirectory(
-      ${CMAKE_CURRENT_SOURCE_DIR}/../../../third-party/ao/torchao/experimental/ops/mps
-      ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/ao/torchao/experimental/ops/mps)
-    target_link_options_shared_lib(torchao_ops_mps_executorch)
-    list(APPEND link_libraries torchao_ops_mps_executorch)
-  endif()
+  # if(CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
+  #   add_subdirectory(
+  #     ${CMAKE_CURRENT_SOURCE_DIR}/../../../third-party/ao/torchao/experimental/ops/mps
+  #     ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/ao/torchao/experimental/ops/mps)
+  #   target_link_options_shared_lib(torchao_ops_mps_executorch)
+  #   list(APPEND link_libraries torchao_ops_mps_executorch)
+  # endif()
 endif()
 
 set(XNNPACK_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../backends/xnnpack)

examples/models/llama/source_transformation/quantize.py

@@ -98,21 +98,38 @@ def quantize(  # noqa C901
         matches = re.findall(pattern, qmode)
         assert len(matches) == 1, f"Expected 1 match for pattern but got {len(matches)}"
         bitwidth = int(matches[0][0])
-        _load_torchao_aten_lib(libname="libtorchao_ops_aten")
-        from torchao.experimental.quant_api import Int8DynActIntxWeightLinearQuantizer
+        # _load_torchao_aten_lib(libname="libtorchao_ops_aten")
+        # from torchao.experimental.quant_api import Int8DynActIntxWeightLinearQuantizer
+        from torchao.experimental.quant_api import int8_dynamic_activation_intx_weight, Int8DynActIntxWeightLinearQuantizer
+        from torchao.quantization.quant_api import quantize_
+        from torchao.utils import unwrap_tensor_subclass
+        from torchao.quantization.granularity import PerRow, PerGroup
 
         with torch.no_grad():
-            model = Int8DynActIntxWeightLinearQuantizer(
-                device="cpu",
-                precision=torch.float32,
-                groupsize=group_size,
-                bitwidth=bitwidth,
-                has_weight_zeros=False,
-            ).quantize(model)
-
+            # model = Int8DynActIntxWeightLinearQuantizer(
+            #     device="cpu",
+            #     precision=torch.float32,
+            #     groupsize=group_size,
+            #     bitwidth=bitwidth,
+            #     has_weight_zeros=False,  
+            # ).quantize(model)
+
+            quantize_(model,
+                int8_dynamic_activation_intx_weight(
+                    # group_size=group_size,
+                    # nbit=bitwidth,
+                    # has_weight_zeros=False,
+                    weight_dtype=getattr(torch, f"int{bitwidth}"),
+                    granularity=PerRow() if group_size == 0 else PerGroup(group_size),
+                    has_weight_zeros=False,
+                ),
+            )
+            model = unwrap_tensor_subclass(model)
         if verbose:
             print("quantized model:", model)
         return model
+
+        return model
     elif qmode == "8da4w":
         # Check for required args
         if group_size is None:
@@ -752,7 +769,7 @@ def get_quant_embedding_transform(args):
         bitwidth, group_size = args.embedding_quantize.split(":")[1].split(",")
         group_size = int(group_size)
         bitwidth = int(bitwidth)
-        _load_torchao_aten_lib(libname="libtorchao_ops_aten")
+        # _load_torchao_aten_lib(libname="libtorchao_ops_aten")
         from torchao.experimental.quant_api import IntxWeightEmbeddingQuantizer
 
         def _torchao_embedding_quantizer(model):

run.sh

@@ -0,0 +1,3 @@
+for i in {1..5}; do
+    ./cmake-out/examples/models/llama/llama_main --model_path=$MODEL_OUT --tokenizer_path=$TOKENIZER --prompt="Once upon a time,"
+done