Qualcomm AI Engine Direct - [LLM QAT] LLM Quant-Aware Distillation (QAD)

backends/qualcomm/_passes/seq_mse.py

@@ -160,6 +160,9 @@ def _insert_seq_mse(
                 # extract observer
                 weight_node_obs = node.args[1]
                 observer = getattr(graph_module, weight_node_obs.name)
+                # For fake quant case
+                if hasattr(observer, "activation_post_process"):
+                    observer = observer.activation_post_process
                 # extract parameters
                 weight_node = weight_node_obs.args[0]
                 weight_tensor = graph_module.get_parameter(weight_node.target).detach()

backends/qualcomm/quantizer/qconfig.py

@@ -896,13 +896,13 @@ def get_16a4w_qnn_qat_config(
         observer_or_fake_quant_ctr=act_fake_quant_ctr,
     )
 
-    weight_fake_quant_ctr = FusedMovingAvgObsFakeQuantize.with_args(
+    weight_fake_quant_ctr = FakeQuantize.with_args(
         dtype=torch.int8,
         quant_min=-7,
         quant_max=7,
         qscheme=torch.per_tensor_symmetric,
         ch_axis=0,
-        observer=MovingAverageMinMaxObserver.with_args(**extra_args),
+        observer=MinMaxObserver.with_args(**extra_args),
     )
     weight_quantization_spec = QuantizationSpec(
         dtype=torch.int8,
@@ -918,7 +918,7 @@ def get_16a4w_qnn_qat_config(
         quant_min=torch.iinfo(torch.int32).min,
         quant_max=torch.iinfo(torch.int32).max,
         qscheme=torch.per_tensor_symmetric,
-        observer=MovingAverageMinMaxObserver.with_args(**extra_args),
+        observer=MinMaxObserver.with_args(**extra_args),
     )
     bias_quantization_spec = QuantizationSpec(
         dtype=torch.int32,
@@ -998,15 +998,15 @@ def get_qat_per_channel_quant_config(
             observer_or_fake_quant_ctr=act_fake_quant_ctr,
         )
 
-    weight_fake_quant_ctr = FusedMovingAvgObsFakeQuantize.with_args(
+    weight_fake_quant_ctr = FakeQuantize.with_args(
         dtype=torch.int8 if weight_dtype == torch.int4 else weight_dtype,
         quant_min=(
             -7 if weight_dtype == torch.int4 else torch.iinfo(weight_dtype).min + 1
         ),
         quant_max=7 if weight_dtype == torch.int4 else torch.iinfo(weight_dtype).max,
         qscheme=torch.per_channel_symmetric,
         ch_axis=ch_axis,
-        observer=MovingAveragePerChannelMinMaxObserver.with_args(**extra_args),
+        observer=PerChannelParamObserver.with_args(**extra_args),
     )
     weight_quantization_spec = QuantizationSpec(
         dtype=torch.int8 if weight_dtype == torch.int4 else weight_dtype,

backends/qualcomm/quantizer/quantizer.py

@@ -207,6 +207,15 @@ class QuantDtype(IntEnum):
         partial(get_qat_per_channel_quant_config),
         None,
     ),
+    (QuantDtype.use_16a8w, True): (
+        get_16a8w_qnn_qat_config,
+        partial(
+            get_qat_per_channel_quant_config,
+            act_dtype=torch.uint16,
+            weight_dtype=torch.int8,
+        ),
+        None,
+    ),
 }
 
 

backends/qualcomm/tests/test_qnn_delegate.py

@@ -8767,6 +8767,80 @@ def test_hf_causal_lm(self):
                         f"Expected Output: '{golden_start_with}' Actual Output: '{model_out}'",
                     )
 
+    def test_static_llm_qat(self):
+        if not self.required_envs():
+            self.skipTest("missing required envs")
+        if self.compile_only:
+            self.skipTest("tasks_eval requires on-device inference")
+
+        def run_eval(
+            calib_limit: int, train_limit: int, extra_args: List[str] = None
+        ) -> float:
+            prompt = "I would like to learn python, could you teach me with a simple example?"
+            cmds = [
+                "python",
+                f"{self.executorch_root}/examples/qualcomm/oss_scripts/llama/llama.py",
+                "--artifact",
+                self.artifact_dir,
+                "--build_folder",
+                self.build_folder,
+                "--prompt",
+                prompt,
+                "--temperature",
+                "0",
+                "--decoder_model",
+                "smollm2_135m",
+                "--model_mode",
+                "kv",
+                "--max_seq_len",
+                "1024",
+                "--max_context_len",
+                "1024",
+                "--eval_methods",
+                "tasks_eval",
+                "--eval_tasks",
+                "wikitext",
+                "--eval_limit",
+                "1",
+                "--qat",
+                "--calib_tasks",
+                "wikitext",
+                "--calib_limit",
+                str(calib_limit),
+                "--train_tasks",
+                "wikitext",
+                "--train_limit",
+                str(train_limit),
+                "--kd_alpha",
+                "0.5",
+            ]
+            if extra_args:
+                cmds.extend(extra_args)
+            self.add_default_cmds(cmds)
+
+            p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL)
+            with Listener((self.ip, self.port)) as listener:
+                conn = listener.accept()
+                p.communicate()
+                msg = json.loads(conn.recv())
+            if "Error" in msg:
+                self.fail(
+                    f"smollm2_135m QAT (limit={train_limit}) failed: {msg['Error']}"
+                )
+            return msg["wiki_ppl"]
+
+        ptq_ppl = run_eval(
+            calib_limit=1, train_limit=1, extra_args=["--freeze_all_params"]
+        )
+        qat_ppl = run_eval(calib_limit=1, train_limit=1)
+        logging.info(f"QAT PPL={qat_ppl:.2f}")
+        logging.info(f"PTQ PPL={ptq_ppl:.2f}")
+        self.assertLess(
+            qat_ppl,
+            ptq_ppl,
+            f"Expected QAT PPL ({qat_ppl:.2f}) < PTQ PPL({ptq_ppl:.2f})",
+        )
+
 
 class TestExampleMultimodalityScript(TestQNN):
 

examples/qualcomm/oss_scripts/llama/README.md

@@ -577,8 +577,6 @@ Calibration data is required for compilation. There are two ways to supply it:
 
 For LLMs, provide at least one of the two; for multimodal models, `--calib_samples` is mandatory.
 
-Calibration and runtime evaluation use separate flag sets and can target different tasks or limits as needed:
-
 | Purpose | Flags |
 |---|---|
 | Calibration data (lm_eval tasks) | `--calib_tasks`, `--calib_limit`, `--calib_num_fewshot` |
@@ -610,19 +608,7 @@ Ready-to-use examples for each model type are provided under `assets/samples/`:
 | ALM (audio) | [assets/samples/audio.json](assets/samples/audio.json) |
 | VLM (vision) | [assets/samples/vision.json](assets/samples/vision.json) |
 
-#### Quantization Guidance
-
-To automatically identify sensitive layers and generate a mixed-precision recipe suggestion, add the `--quant_recipe_suggestion` flag. During calibration, the analyzer compares FP32 and QDQ intermediate outputs layer-by-layer using SQNR, then writes two files to the working directory:
-
-- `{model_name}_quantization_error.csv` — per-group SQNR statistics sorted by sensitivity (most sensitive first)
-- `{model_name}_suggest_recipe.py` — ready-to-use `StaticLLMQuantRecipe` subclasses optimized to apply higher-precision quantization to the most sensitive groups.
-
-Example:
-```bash
-python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} --prompt "I would like to learn python, could you teach me with a simple example?" --temperature 0 --model_mode kv --max_seq_len 1024 --decoder_model qwen3-1_7b --calib_tasks wikitext --calib_limit 1 --quant_recipe_suggestion --compile_only
-```
-
-After the run, pick one of the generated classes from `qwen3-1_7b_suggest_recipe.py` as your new recipe. For a full walkthrough, see [quantization_guidance.md](quantization_guidance.md).
+For the full quantization workflow: PTQ (Post Training Quantization), QAT(Quantization Aware Training / Distillation), Mixed-precision, Quantization error analysis, see [quantization_guidance.md](quantization_guidance.md).
 
 #### Use attention sink for multi-turn conversations
 Attention sink is a way to evict cache when maximum context length be reached.

examples/qualcomm/oss_scripts/llama/__init__.py

@@ -83,6 +83,7 @@
     Qwen2_5_1_5BQuantRecipe,
     Qwen3_0_6BQuantRecipe,
     Qwen3_1_7BQuantRecipe,
+    Smollm2QATQuantRecipe,
     Smollm2QuantRecipe,
     Smollm3QuantRecipe,
     SmolVLMQuantRecipe,
@@ -125,6 +126,8 @@ class LLMModelConfig(ABC):
     r2: Enable SpinQuant R2 quantization optimization.
     r3: Enable SpinQuant R3 quantization optimization.
     quant_recipe: Quantization recipe to use when setting quant configs for the model.
+    qat_recipe: Optional QAT-specific quantization recipe. Used only when --qat is set;
+                if None, falls back to quant_recipe.
     """
 
     repo_id: str
@@ -141,6 +144,7 @@ class LLMModelConfig(ABC):
     r2: bool
     r3: bool
     quant_recipe: StaticLLMQuantRecipe
+    qat_recipe: Optional[Type[StaticLLMQuantRecipe]] = None
 
     def __str__(self):  # noqa: C901
         """
@@ -541,6 +545,7 @@ class Smollm2_135M(LLMModelConfig):
     r2 = False
     r3 = True
     quant_recipe = Smollm2QuantRecipe
+    qat_recipe = Smollm2QATQuantRecipe
 
 
 @register_llm_model("smollm3-3b")

examples/qualcomm/oss_scripts/llama/dataset/__init__.py

@@ -11,8 +11,12 @@
 )
 from executorch.examples.qualcomm.oss_scripts.llama.dataset.collators import (
     LLMCalibCollator,
+    LLMTrainingCollator,
 )
 from executorch.examples.qualcomm.oss_scripts.llama.dataset.config import DataConfig
+from executorch.examples.qualcomm.oss_scripts.llama.dataset.constants import (
+    LABEL_IGNORE_INDEX,
+)
 from executorch.examples.qualcomm.oss_scripts.llama.dataset.datasets import (
     LLMDataset,
     ModalityEncoderDataset,
@@ -21,28 +25,40 @@
     collect_lm_eval_tokens,
     load_audio_file,
     load_conversation_samples,
+    load_hf_chat_dataset,
 )
 from executorch.examples.qualcomm.oss_scripts.llama.dataset.preprocessors import (
     ModalityPreprocessor,
     preprocess_encoder_inputs,
 )
 from executorch.examples.qualcomm.oss_scripts.llama.dataset.schema import MessageSample
+from executorch.examples.qualcomm.oss_scripts.llama.dataset.targets import (
+    make_causal_labels,
+    make_conversation_labels,
+)
 
 __all__ = [
     # config
     "DataConfig",
     # schema
     "MessageSample",
+    # constants
+    "LABEL_IGNORE_INDEX",
     # loaders
     "collect_lm_eval_tokens",
     "load_audio_file",
     "load_conversation_samples",
+    "load_hf_chat_dataset",
+    # targets (label generation — low-level functions)
+    "make_causal_labels",
+    "make_conversation_labels",
     # datasets
     "LLMDataset",
     "ModalityEncoderDataset",
     # collators
     "ModalityEncoderCollator",
     "LLMCalibCollator",
+    "LLMTrainingCollator",
     # builders
     "DatasetBuilder",
     "DecoderDatasetBuilder",