Enable kwarg inputs for pt2e quantize for torchtune models

jackzhxng · jackzhxng · commit 945cdf531743 · 2025-01-02T11:37:59.000-08:00
diff --git a/extension/llm/export/builder.py b/extension/llm/export/builder.py
@@ -342,7 +342,7 @@ def pt2e_quantize(self, quantizers: Optional[List[Quantizer]]) -> "LLMEdgeManage
                     logging.info(
                         "No calibration provided, using dummy input to calibrate..."
                     )
-                    m(*self.example_inputs)
+                    m(*self.example_inputs, **self.example_kwarg_inputs)
                 m = convert_pt2e(m)
                 DuplicateDynamicQuantChainPass()(m)
                 self.pre_autograd_graph_module = m

Original file line number	Diff line number	Diff line change
`@@ -342,7 +342,7 @@ def pt2e_quantize(self, quantizers: Optional[List[Quantizer]]) -> "LLMEdgeManage`
`342`	`342`	`logging.info(`
`343`	`343`	`"No calibration provided, using dummy input to calibrate..."`
`344`	`344`	`)`
`345`		`- m(*self.example_inputs)`
	`345`	`+ m(self.example_inputs, *self.example_kwarg_inputs)`
`346`	`346`	`m = convert_pt2e(m)`
`347`	`347`	`DuplicateDynamicQuantChainPass()(m)`
`348`	`348`	`self.pre_autograd_graph_module = m`