From 4655850dcef8bbeba3f75845403b7bc4e6979176 Mon Sep 17 00:00:00 2001
From: Jack <32371937+jackzhxng@users.noreply.github.com>
Date: Thu, 27 Feb 2025 10:07:03 -0800
Subject: [PATCH] Enable kwarg inputs for pt2e quantize (#7436)

Summary:
For quantizing models that have kwarg forward() inputs, such as TorchTune Llama models


Test Plan: N/A

Reviewed By: tarun292

Differential Revision: D70206003

Pulled By: jackzhxng
---
 extension/llm/export/builder.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/extension/llm/export/builder.py b/extension/llm/export/builder.py
index 3e987489693..47ad30e9390 100644
--- a/extension/llm/export/builder.py
+++ b/extension/llm/export/builder.py
@@ -360,7 +360,10 @@ def pt2e_quantize(self, quantizers: Optional[List[Quantizer]]) -> "LLMEdgeManage
                     logging.info(
                         "No calibration provided, using dummy input to calibrate..."
                     )
-                    m(*self.example_inputs)
+                    if self.example_kwarg_inputs:
+                        m(*self.example_inputs, **self.example_kwarg_inputs)
+                    else:
+                        m(*self.example_inputs)
                 m = convert_pt2e(m)
                 DuplicateDynamicQuantChainPass()(m)
                 self.pre_autograd_graph_module = m