Skip to content

Commit c461421

Browse files
Qualcomm AI Engine Direct - [LLM QAT] LLM Quant-Aware Distillation (QAD) (pytorch#21036)
### Summary Training: - Add BaseTrainer, Trainer (CE), KDTrainer (knowledge distillation) - Add CrossEntropyLoss, KLDivergenceLoss, linear warm-up cosine LR scheduler - Add TrainingArgs dataclass (epochs, lr, alpha, temperature, grad_accum_steps, warmup_ratio, lr_config) Data pipeline: - Add DecoderDatasetBuilder.from_hf_source for HuggingFace SFT chat datasets - Add build_qat_dataloaders: explicit calib/train split - Add LLMTrainingCollator, make_causal_labels, make_conversation_labels - Add DataConfig train fields: train_tasks, train_hf_dataset, train_hf_limit Quantization Strategy: - Add QATStrategy: PTQ calibration pass followed by KDTrainer/Trainer fine-tuning - Branch TextDecoder.quantize on --qat: prepare_qat_pt2e + move_exported_model_to_train vs prepare_pt2e - Select qat_recipe over quant_recipe when --qat is active Quant recipe: - Add StaticLLMQATRecipe base class as a example - Add Smollm2QATQuantRecipe; add qat_recipe field to LLMModelConfig - Add 16a8w QAT qconfig Fix: - SeqMSE: unwrap FakeQuantize wrapper before extracting observer CI Testing: - Add test_static_llm_qat: assert QAT PPL < PTQ PPL on smollm2_135m README: - `examples/qualcomm/oss_scripts/llama/quantization_guidance.md` #### E2E script: ``` bash python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} --decoder_model smollm2_135m --model_mode hybrid --max_seq_len 1024 --prompt "What dose it mean to edit written content?" --qat --train_hf_dataset "HuggingFaceTB/smol-smoltalk" --train_hf_limit 1000 --calib_hf_dataset "HuggingFaceTB/smol-smoltalk" --calib_hf_limit 1000 ``` #### Result ``` bash I 00:00:00.284922 executorch:runner.cpp:228] Reading metadata from model I 00:00:00.285219 executorch:runner.cpp:360] creating io_memory I 00:00:00.285856 executorch:prompt_processor.cpp:273] Prompt Processor: total 18 prompt tokens (AR-32 * 1 iters) I 00:00:00.295011 executorch:runner.cpp:473] RSS after prompt prefill: 181.175781 MiB (0 if unsupported) I 00:00:01.262212 executorch:token_generator.cpp:347] Reached to the end of generation I 00:00:01.262249 executorch:runner.cpp:490] RSS after finishing text generation: 181.175781 MiB (0 if unsupported) I 00:00:01.262275 executorch:stats.h:161] Prompt Tokens: 18 Generated Tokens: 355 I 00:00:01.262284 executorch:stats.h:167] Model Load Time: 0.284000 (seconds) I 00:00:01.262291 executorch:stats.h:176] Total inference time: 0.977000 (seconds) Rate: 363.357216 (tokens/second) I 00:00:01.262298 executorch:stats.h:185] Prompt evaluation: 0.010000 (seconds) Rate: 1800.000000 (tokens/second) I 00:00:01.262305 executorch:stats.h:197] Generated 355 tokens: 0.967000 (seconds) Rate: 367.114788 (tokens/second) I 00:00:01.262311 executorch:stats.h:205] Time to first generated token: 0.010000 (seconds) I 00:00:01.262318 executorch:stats.h:212] Sampling time over 373 tokens: 0.112000 (seconds) [INFO] [Qnn ExecuTorch]: Destroy Qnn context [INFO] [Qnn ExecuTorch]: Destroy Qnn device [INFO] [Qnn ExecuTorch]: Destroy Qnn backend PyTorchObserver {"prefill_token_per_sec":1800,"decode_token_per_sec":367.115,"prompt_tokens":18,"generated_tokens":355,"model_load_start_ms":1753685358963,"model_load_end_ms":1753685359247,"inference_start_ms":1753685359247,"inference_end_ms":1753685360224,"prompt_eval_end_ms":1753685359257,"first_token_ms":1753685359257,"aggregate_sampling_time_ms":112,"SCALING_FACTOR_UNITS_PER_SECOND":1000} /data/local/tmp/yuyazhua/executorch/static_llm/outputs/outputs.txt: 1 file pulled, 0 skipped. 1.8 MB/s (1714 bytes in 0.001s) /data/local/tmp/yuyazhua/executorch/static_llm/outputs/inference_speed.txt: 1 file pulled, 0 skipped. 0.0 MB/s (7 bytes in 0.001s) [INFO 2026-07-20 12:10:52,191 llama.py:290] Device Inference Results[0]: <|im_start|>user What dose it mean to edit written content?<|im_end|> <|im_start|>assistant 1. **Editing for clarity**: When you edit your writing, you aim to make it clear and concise, ensuring that your message is easy for readers to understand and retain. This involves cutting unnecessary words, phrases, or sentences, and rephrasing or reorganizing your content to make it more readable. 2. **Editing for grammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errors in grammar, syntax, and punctuation, which can make your writing more polished and error-free. 3. **Editing for style**: The final stage of editing is where you refine your writing to make it more engaging, clear, and engaging. This involves using language that is engaging, yet clear, and engaging in the way you want it to be. 4. **Editing for impact**: The final stage of editing is where you aim to make your writing more impactful. This means that you strive to make your message clear, concise, and impactful, and to convey it in a way that resonates with your readers. 5. **Editing for consistency**: When you edit your writing, you ensure that it is consistent in terms of tone, style, and voice. This involves using the same language, vocabulary, and structure to create a consistent tone, which is essential for effective communication. To get started, take a step back and look at your writing. Read it aloud, if possible, to get a sense of the tone, style, and voice you're using. Then, read your writing aloud to get a sense of the rhythm, cadence, and flow. When you're on the right track, you can begin to refine your writing and make it more engaging and effective.<|im_end|> ``` ### Test plan ``` bash python backends/qualcomm/tests/test_qnn_delegate.py TestExampleLLMScript.test_static_llm_qat --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} -r . -a . --build_folder build-android ``` cc: @shewu-quic @haowhsu-quic @winskuo-quic @psiddh @abhinaykukkadapu
1 parent 4753d75 commit c461421

35 files changed

Lines changed: 1870 additions & 80 deletions

‎backends/qualcomm/_passes/seq_mse.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,9 @@ def _insert_seq_mse(
160160
# extract observer
161161
weight_node_obs = node.args[1]
162162
observer = getattr(graph_module, weight_node_obs.name)
163+
# For fake quant case
164+
if hasattr(observer, "activation_post_process"):
165+
observer = observer.activation_post_process
163166
# extract parameters
164167
weight_node = weight_node_obs.args[0]
165168
weight_tensor = graph_module.get_parameter(weight_node.target).detach()

‎backends/qualcomm/quantizer/qconfig.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -998,15 +998,15 @@ def get_qat_per_channel_quant_config(
998998
observer_or_fake_quant_ctr=act_fake_quant_ctr,
999999
)
10001000

1001-
weight_fake_quant_ctr = FusedMovingAvgObsFakeQuantize.with_args(
1001+
weight_fake_quant_ctr = FakeQuantize.with_args(
10021002
dtype=torch.int8 if weight_dtype == torch.int4 else weight_dtype,
10031003
quant_min=(
10041004
-7 if weight_dtype == torch.int4 else torch.iinfo(weight_dtype).min + 1
10051005
),
10061006
quant_max=7 if weight_dtype == torch.int4 else torch.iinfo(weight_dtype).max,
10071007
qscheme=torch.per_channel_symmetric,
10081008
ch_axis=ch_axis,
1009-
observer=MovingAveragePerChannelMinMaxObserver.with_args(**extra_args),
1009+
observer=PerChannelParamObserver.with_args(**extra_args),
10101010
)
10111011
weight_quantization_spec = QuantizationSpec(
10121012
dtype=torch.int8 if weight_dtype == torch.int4 else weight_dtype,

‎backends/qualcomm/quantizer/quantizer.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,15 @@ class QuantDtype(IntEnum):
207207
partial(get_qat_per_channel_quant_config),
208208
None,
209209
),
210+
(QuantDtype.use_16a8w, True): (
211+
get_16a8w_qnn_qat_config,
212+
partial(
213+
get_qat_per_channel_quant_config,
214+
act_dtype=torch.uint16,
215+
weight_dtype=torch.int8,
216+
),
217+
None,
218+
),
210219
}
211220

212221

‎backends/qualcomm/tests/test_qnn_delegate.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9149,6 +9149,78 @@ def test_hf_causal_lm(self):
91499149
f"Expected Output: '{golden_start_with}' Actual Output: '{model_out}'",
91509150
)
91519151

9152+
def test_static_llm_qat(self):
9153+
if not self.required_envs():
9154+
self.skipTest("missing required envs")
9155+
if self.compile_only:
9156+
self.skipTest("tasks_eval requires on-device inference")
9157+
9158+
def run_eval(
9159+
calib_limit: int, train_limit: int, extra_args: List[str] = None
9160+
) -> float:
9161+
prompt = "I would like to learn python, could you teach me with a simple example?"
9162+
cmds = [
9163+
"python",
9164+
f"{self.executorch_root}/examples/qualcomm/oss_scripts/llama/llama.py",
9165+
"--artifact",
9166+
self.artifact_dir,
9167+
"--build_folder",
9168+
self.build_folder,
9169+
"--prompt",
9170+
prompt,
9171+
"--temperature",
9172+
"0",
9173+
"--decoder_model",
9174+
"smollm2_135m",
9175+
"--model_mode",
9176+
"kv",
9177+
"--max_seq_len",
9178+
"1024",
9179+
"--max_context_len",
9180+
"1024",
9181+
"--eval_methods",
9182+
"tasks_eval",
9183+
"--eval_tasks",
9184+
"wikitext",
9185+
"--eval_limit",
9186+
"1",
9187+
"--qat",
9188+
"--calib_tasks",
9189+
"wikitext",
9190+
"--calib_limit",
9191+
str(calib_limit),
9192+
"--train_tasks",
9193+
"wikitext",
9194+
"--train_limit",
9195+
str(train_limit),
9196+
]
9197+
if extra_args:
9198+
cmds.extend(extra_args)
9199+
self.add_default_cmds(cmds)
9200+
9201+
p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL)
9202+
with Listener((self.ip, self.port)) as listener:
9203+
conn = listener.accept()
9204+
p.communicate()
9205+
msg = json.loads(conn.recv())
9206+
if "Error" in msg:
9207+
self.fail(
9208+
f"smollm2_135m QAT (limit={train_limit}) failed: {msg['Error']}"
9209+
)
9210+
return msg["wiki_ppl"]
9211+
9212+
ptq_ppl = run_eval(
9213+
calib_limit=1, train_limit=1, extra_args=["--freeze_all_params"]
9214+
)
9215+
qat_ppl = run_eval(calib_limit=1, train_limit=1)
9216+
logging.info(f"QAT PPL={qat_ppl:.2f}")
9217+
logging.info(f"PTQ PPL={ptq_ppl:.2f}")
9218+
self.assertLess(
9219+
qat_ppl,
9220+
ptq_ppl,
9221+
f"Expected QAT PPL ({qat_ppl:.2f}) < PTQ PPL({ptq_ppl:.2f})",
9222+
)
9223+
91529224

91539225
class TestExampleMultimodalityScript(TestQNN):
91549226

‎examples/qualcomm/oss_scripts/llama/README.md‎

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -584,8 +584,6 @@ Calibration data is required for compilation. There are two ways to supply it:
584584

585585
For LLMs, provide at least one of the two; for multimodal models, `--calib_samples` is mandatory.
586586

587-
Calibration and runtime evaluation use separate flag sets and can target different tasks or limits as needed:
588-
589587
| Purpose | Flags |
590588
|---|---|
591589
| Calibration data (lm_eval tasks) | `--calib_tasks`, `--calib_limit`, `--calib_num_fewshot` |
@@ -617,19 +615,7 @@ Ready-to-use examples for each model type are provided under `assets/samples/`:
617615
| ALM (audio) | [assets/samples/audio.json](assets/samples/audio.json) |
618616
| VLM (vision) | [assets/samples/vision.json](assets/samples/vision.json) |
619617

620-
#### Quantization Guidance
621-
622-
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:
623-
624-
- `{model_name}_quantization_error.csv` — per-group SQNR statistics sorted by sensitivity (most sensitive first)
625-
- `{model_name}_suggest_recipe.py` — ready-to-use `StaticLLMQuantRecipe` subclasses optimized to apply higher-precision quantization to the most sensitive groups.
626-
627-
Example:
628-
```bash
629-
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
630-
```
631-
632-
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).
618+
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).
633619

634620
#### Use attention sink for multi-turn conversations
635621
Attention sink is a way to evict cache when maximum context length be reached.

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@
8787
Qwen2_5_1_5BQuantRecipe,
8888
Qwen3_0_6BQuantRecipe,
8989
Qwen3_1_7BQuantRecipe,
90+
Smollm2QATQuantRecipe,
9091
Smollm2QuantRecipe,
9192
Smollm3QuantRecipe,
9293
SmolVLMQuantRecipe,
@@ -129,6 +130,8 @@ class LLMModelConfig(ABC):
129130
r2: Enable SpinQuant R2 quantization optimization.
130131
r3: Enable SpinQuant R3 quantization optimization.
131132
quant_recipe: Quantization recipe to use when setting quant configs for the model.
133+
qat_recipe: Optional QAT-specific quantization recipe. Used only when --qat is set;
134+
if None, falls back to quant_recipe.
132135
"""
133136

134137
repo_id: str
@@ -145,6 +148,7 @@ class LLMModelConfig(ABC):
145148
r2: bool
146149
r3: bool
147150
quant_recipe: StaticLLMQuantRecipe
151+
qat_recipe: Optional[Type[StaticLLMQuantRecipe]] = None
148152

149153
def __str__(self): # noqa: C901
150154
"""
@@ -545,6 +549,7 @@ class Smollm2_135M(LLMModelConfig):
545549
r2 = False
546550
r3 = True
547551
quant_recipe = Smollm2QuantRecipe
552+
qat_recipe = Smollm2QATQuantRecipe
548553

549554

550555
@register_llm_model("smollm3-3b")
114 KB
Loading
140 KB
Loading
119 KB
Loading

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,12 @@
1111
)
1212
from executorch.examples.qualcomm.oss_scripts.llama.dataset.collators import (
1313
LLMCalibCollator,
14+
LLMTrainingCollator,
1415
)
1516
from executorch.examples.qualcomm.oss_scripts.llama.dataset.config import DataConfig
17+
from executorch.examples.qualcomm.oss_scripts.llama.dataset.constants import (
18+
LABEL_IGNORE_INDEX,
19+
)
1620
from executorch.examples.qualcomm.oss_scripts.llama.dataset.datasets import (
1721
LLMDataset,
1822
ModalityEncoderDataset,
@@ -21,28 +25,40 @@
2125
collect_lm_eval_tokens,
2226
load_audio_file,
2327
load_conversation_samples,
28+
load_hf_chat_dataset,
2429
)
2530
from executorch.examples.qualcomm.oss_scripts.llama.dataset.preprocessors import (
2631
ModalityPreprocessor,
2732
preprocess_encoder_inputs,
2833
)
2934
from executorch.examples.qualcomm.oss_scripts.llama.dataset.schema import MessageSample
35+
from executorch.examples.qualcomm.oss_scripts.llama.dataset.targets import (
36+
make_causal_labels,
37+
make_conversation_labels,
38+
)
3039

3140
__all__ = [
3241
# config
3342
"DataConfig",
3443
# schema
3544
"MessageSample",
45+
# constants
46+
"LABEL_IGNORE_INDEX",
3647
# loaders
3748
"collect_lm_eval_tokens",
3849
"load_audio_file",
3950
"load_conversation_samples",
51+
"load_hf_chat_dataset",
52+
# targets (label generation — low-level functions)
53+
"make_causal_labels",
54+
"make_conversation_labels",
4055
# datasets
4156
"LLMDataset",
4257
"ModalityEncoderDataset",
4358
# collators
4459
"ModalityEncoderCollator",
4560
"LLMCalibCollator",
61+
"LLMTrainingCollator",
4662
# builders
4763
"DatasetBuilder",
4864
"DecoderDatasetBuilder",

0 commit comments

Comments
 (0)