up

Author: metascroy

.ci/scripts/test_torchao_huggingface_checkpoints.sh

@@ -51,7 +51,7 @@ MODEL_OUT=model.pte
 # Default to XNNPACK
 BACKEND_ARGS="-X --xnnpack-extended-ops"
 if [[ "$USE_TORCHAO_KERNELS" -eq 1 ]]; then
-  BACKEND_ARGS="--torchao-kernels"
+  BACKEND_ARGS="--use-torchao-kernels"
 fi
 
 case "$MODEL_NAME" in

examples/models/llama/export_llama_lib.py

@@ -418,17 +418,17 @@ def build_args_parser() -> argparse.ArgumentParser:
         help="Delegate more operators beyond DQLinear to the xnnpack backend. Requires -X or --xnnpack to be set.",
     )
     parser.add_argument(
-        "--torchao-kernels",
+        "--use-torchao-kernels",
         action="store_true",
         help="Delegate tied-embedding and quantized linear ops to torchao kernels",
     )
     parser.add_argument(
-        "--torchao-kernels-tied-embedding",
+        "--use-torchao-kernels-tied-embedding",
         action="store_true",
         help="Delegate tied-embedding ops to torchao kernels",
     )
     parser.add_argument(
-        "--torchao-kernels-linear",
+        "--use-torchao-kernels-linear",
         action="store_true",
         help="Delegate linear ops to torchao kernels",
     )
@@ -756,7 +756,7 @@ def _prepare_for_llama_export(llm_config: LlmConfig) -> LLMEdgeManager:
             preq_group_size=llm_config.base.preq_group_size,
             preq_embedding_quantize=llm_config.base.preq_embedding_quantize,
             local_global_attention=llm_config.model.local_global_attention,
-            use_torchao_kernels_linear=llm_config.backend.torchao.convert_linear,
+            use_torchao_kernels_linear=llm_config.backend.torchao.use_torchao_kernels_linear,
             use_torchao_kernels_tied_embedding=llm_config.backend.torchao.convert_tied_embedding,
         )
     )

extension/llm/export/config/llm_config.py

@@ -458,8 +458,8 @@ class TorchAOKernelsConfig:
     Configures the torchao-kernels backend.
     """
 
-    convert_linear: bool = False
-    convert_tied_embedding: bool = False
+    use_torchao_kernels_linear: bool = False
+    use_torchao_kernels_tied_embedding: bool = False
 
 
 @dataclass
@@ -647,22 +647,22 @@ def from_args(cls, args: argparse.Namespace) -> "LlmConfig":  # noqa: C901
         if any(
             hasattr(args, a)
             for a in [
-                "torchao_kernels",
-                "torchao_kernels_linear",
-                "torchao_kernels_tied_embedding",
+                "use_torchao_kernels",
+                "use_torchao_kernels_linear",
+                "use_torchao_kernels_tied_embedding",
             ]
         ):
-            if hasattr(args, "torchao_kernels") and args.torchao_kernels:
+            if hasattr(args, "use_torchao_kernels") and args.torchao_kernels:
                 # Enable all conversions if torchao_kernels is specified
-                llm_config.backend.torchao.convert_linear = True
-                llm_config.backend.torchao.convert_tied_embedding = True
+                llm_config.backend.torchao.use_torchao_kernels_linear = True
+                llm_config.backend.torchao.use_torchao_kernels_tied_embedding = True
             else:
                 # Otherwise, only enable the conversions that are specified
-                llm_config.backend.torchao.convert_linear = getattr(
-                    args, "torchao_kernels_linear", False
+                llm_config.backend.torchao.use_torchao_kernels_linear = getattr(
+                    args, "use_torchao_kernels_linear", False
                 )
-                llm_config.backend.torchao.convert_tied_embedding = getattr(
-                    args, "torchao_kernels_tied_embedding", False
+                llm_config.backend.torchao.use_torchao_kernels_tied_embedding = getattr(
+                    args, "use_torchao_kernels_tied_embedding", False
                 )
 
         # DebugConfig