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

Merged
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD
Aug 13, 2026
Merged

Qualcomm AI Engine Direct - [LLM QAT] LLM Quant-Aware Distillation (QAD)#21036
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD

Conversation

@DannyYuyang-quic

@DannyYuyang-quicDannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

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:

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

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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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

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

@pytorch-bot

pytorch-botBot commented Jul 20, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21036

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job

As of commit 54a691b with merge base bc2833f (image):

CANCELLED JOB - The following job was cancelled. Please retry:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 20, 2026
@DannyYuyang-quic

DannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
ContributorAuthor

@psiddh Hi,
As more users experiment with lower-bit quantization, maintaining model quality with PTQ alone becomes increasingly challenging. To improve output quality under more aggressive quantization settings, we've explored and added a simple end-to-end QAT (Quantization-Aware Training) flow for LLMs.

With QAT, we can further push quantization down to W4 PCQ, and potentially even lower precisions (W2). Below is a comparison between PTQ and QAT on SmolLM2-135M-instruct using W4 per-channel quantization for all decoder layers.

Please have a look, thanks!

Experiment Results: W4 Decoder PCQ (QAT vs PTQ)

QAT (W4 PCQ decoder)

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

[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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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|>

PTQ (W4 PCQ decoder)

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?" --calib_hf_dataset "HuggingFaceTB/smol-smoltalk" --calib_hf_limit 1000

Result

[INFO 2026-07-20 15:12:57,609 llama.py:290] Device Inference Results[0]:
<|im_start|>user
What dose it mean to edit written content?<|im_end|><|im_start|>assistant
1. A sentence is usually written in the first person, i.e., "He said, she said" or "She said, she said."
2. A paragraph is usually written in the third person, i.e., "She said, he said" or "She said, she said"
3. A couple is usually written in the second person, i.e., "She said, he said" or "He said, she said"
4. A little is usually written in the third person, i.e., "She said, he said" or "She said, she said"
5. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
6. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
7. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
8. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
9. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
10. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "He said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "

@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@pytorchbot label "release notes: qualcomm"

@pytorch-botpytorch-botBot added the release notes: qualcomm Changes to the Qualcomm backend delegate label Jul 20, 2026
@psiddh

Copy link
Copy Markdown
Contributor

@claude review

@psiddh

Copy link
Copy Markdown
Contributor

@claude

claudeBot commented Aug 11, 2026

Copy link
Copy Markdown

Claude finished @psiddh's task in 5m 52s —— View job


Claude Code Review

Reviewed the full diff (git diff origin/main...HEAD, 35 files). This is a well-structured QAT/QAD feature: the trainer abstraction, injected forward_fn, and PTQ-then-train strategy composition are clean and the CI test that asserts QAT PPL < PTQ PPL is a nice regression guard. Below are the issues I found, ordered by severity.


🔴 Blocking

1. CI test passes --kd_alpha, an argument that doesn't exist
backends/qualcomm/tests/test_qnn_delegate.py:8814 appends "--kd_alpha", "0.5" to the command, but there is no --kd_alpha in _build_parser() (the KD weight is alpha, sourced only from the --train_config YAML). Since llama.py:1022 uses strict parser.parse_args(), the subprocess will exit with error: unrecognized arguments: --kd_alpha 0.5, and the test's socket Listener will hang or fail rather than report a clean error. This test cannot currently pass as written.

Either register a --kd_alpha arg that overrides TrainingArgs.alpha, or drop it from the test command. Fix this →


🟠 Correctness

2. make_conversation_labels can write to labels[-1] via negative indexing
dataset/targets.py:40-42:

foriinrange(n-1):
ifi<len(assistant_mask) andassistant_mask[i]:
labels[i-1] =tokens[i]

When i == 0 and assistant_mask[0] is truthy, labels[-1] = tokens[0] silently writes to the last label position (Python negative index) instead of being a no-op. In practice position 0 is usually BOS/system so the mask is 0 there, but this is a latent data-corruption bug — one mislabeled sample can inject a spurious supervised target at the final position. Guard with if i > 0: (or start the range at 1). Fix this →

3. Backend qconfig change affects existing QAT callers, not just this feature
backends/qualcomm/quantizer/qconfig.py swaps FusedMovingAvgObsFakeQuantizeFakeQuantize and MovingAverage*ObserverMinMaxObserver/PerChannelParamObserver in get_16a4w_qnn_qat_config and get_qat_per_channel_quant_config. These functions are also used by annotators/htp_rules.py and annotators/lpai_rules.py. Switching from moving-average to plain min/max observers changes calibration behavior for all existing consumers of these QAT configs. This aligns with the new strategy (observers are frozen after the calibration pass via _enable_fake_quant_freeze_observer), but it's an implicit behavior change outside the LLM example scope. Please confirm this is intended for the htp/lpai rule paths and call it out in the PR summary.


🟡 Minor / Questions

4. Operator-precedence readability in the use_qat guardllm_wrappers.py:599-603:

if (
self.mode==Mode.CALIBRATEandself.control_args.quant_recipe_suggestionoruse_qat
):

This parses as (mode==CALIBRATE and quant_recipe_suggestion) or use_qat, which is correct (use_qat already implies CALIBRATE), but the mixed and/or without parens is easy to misread. Wrap the and clause in explicit parentheses.

5. _split_train_val / _split_qat_dataset can produce an all-drop_last=True empty loader. With small --train_limit (the CI uses train_limit=1) and batch_size=1 this is fine, but n_train = max(1, ...) followed by drop_last=True means a training subset with a single sample that doesn't fill batch_size>1 yields zero batches → len(train_loader)==0total_steps==0. _log_dataset_stats guards the calib path but the post-split train subset isn't re-checked. Consider validating len(train_loader) > 0 after the split.

6. KLDivergenceLoss zero-path device but not dtypeloss.py:59 returns torch.zeros(1, device=...).squeeze() (float32) as a fallback; fine on CPU, but if the model is ever autocast/half this won't match student_logits.dtype. Low priority since device=torch.device("cpu") is hardcoded today.

7. BaseTrainer.device is hardcoded to CPU (trainer.py:45, with a # TODO: validate GPU compatibility). Worth noting in the guidance doc that QAT is CPU-only today — for larger decoders this training pass could be very slow.

8. Stale docstring in build_qat_dataloaders references --calib_train_ratio / --train_val_ratio for the "full" mode; the explicit-split branch docstring is accurate, but a reader skimming the method header sees "Full split" first. Minor.


Nits

  • llama.py:528 adds a stray blank line after --r3 default.
  • Help text typos: "lm-eluther"lm-eval / lm-evaluation-harness (appears in --qat_full_tasks and --train_tasks help).
  • train/config/config.py docstring says defaults live in train/config/config.py but the parser help (llama.py:659) points to train/config/config.py too — consistent, good; just confirm the qad.yaml comment "train/config.py" (missing config/) is updated.

Overall the design is solid and the abstractions are clean. Item #1 blocks the test from passing and #2/#3 are the substantive correctness items to resolve before merge.
· branch dev1/danny/llm_QAT_QAD

@psiddh

Copy link
Copy Markdown
Contributor

@DannyYuyang-quic can you resolve teh merge conflicts ?

…for LLMs
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
Testing:
- Add test_static_llm_qat: assert QAT PPL < PTQ PPL on smollm2_135m
@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@psiddh I've rebased the branch and addressed code review comments. Please take a look, thanks!

@psiddh
psiddh merged commit c461421 into pytorch:mainAug 13, 2026
195 of 196 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.release notes: qualcommChanges to the Qualcomm backend delegate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@DannyYuyang-quic@psiddh
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Skip to content

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

Merged
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD
Aug 13, 2026
Merged

Qualcomm AI Engine Direct - [LLM QAT] LLM Quant-Aware Distillation (QAD)#21036
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD

Conversation

@DannyYuyang-quic

@DannyYuyang-quicDannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

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:

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

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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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

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

@pytorch-bot

pytorch-botBot commented Jul 20, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21036

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job

As of commit 54a691b with merge base bc2833f (image):

CANCELLED JOB - The following job was cancelled. Please retry:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 20, 2026
@DannyYuyang-quic

DannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
ContributorAuthor

@psiddh Hi,
As more users experiment with lower-bit quantization, maintaining model quality with PTQ alone becomes increasingly challenging. To improve output quality under more aggressive quantization settings, we've explored and added a simple end-to-end QAT (Quantization-Aware Training) flow for LLMs.

With QAT, we can further push quantization down to W4 PCQ, and potentially even lower precisions (W2). Below is a comparison between PTQ and QAT on SmolLM2-135M-instruct using W4 per-channel quantization for all decoder layers.

Please have a look, thanks!

Experiment Results: W4 Decoder PCQ (QAT vs PTQ)

QAT (W4 PCQ decoder)

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

[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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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|>

PTQ (W4 PCQ decoder)

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?" --calib_hf_dataset "HuggingFaceTB/smol-smoltalk" --calib_hf_limit 1000

Result

[INFO 2026-07-20 15:12:57,609 llama.py:290] Device Inference Results[0]:
<|im_start|>user
What dose it mean to edit written content?<|im_end|><|im_start|>assistant
1. A sentence is usually written in the first person, i.e., "He said, she said" or "She said, she said."
2. A paragraph is usually written in the third person, i.e., "She said, he said" or "She said, she said"
3. A couple is usually written in the second person, i.e., "She said, he said" or "He said, she said"
4. A little is usually written in the third person, i.e., "She said, he said" or "She said, she said"
5. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
6. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
7. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
8. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
9. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
10. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "He said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "

@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@pytorchbot label "release notes: qualcomm"

@pytorch-botpytorch-botBot added the release notes: qualcomm Changes to the Qualcomm backend delegate label Jul 20, 2026
@psiddh

Copy link
Copy Markdown
Contributor

@claude review

@psiddh

Copy link
Copy Markdown
Contributor

@claude

claudeBot commented Aug 11, 2026

Copy link
Copy Markdown

Claude finished @psiddh's task in 5m 52s —— View job


Claude Code Review

Reviewed the full diff (git diff origin/main...HEAD, 35 files). This is a well-structured QAT/QAD feature: the trainer abstraction, injected forward_fn, and PTQ-then-train strategy composition are clean and the CI test that asserts QAT PPL < PTQ PPL is a nice regression guard. Below are the issues I found, ordered by severity.


🔴 Blocking

1. CI test passes --kd_alpha, an argument that doesn't exist
backends/qualcomm/tests/test_qnn_delegate.py:8814 appends "--kd_alpha", "0.5" to the command, but there is no --kd_alpha in _build_parser() (the KD weight is alpha, sourced only from the --train_config YAML). Since llama.py:1022 uses strict parser.parse_args(), the subprocess will exit with error: unrecognized arguments: --kd_alpha 0.5, and the test's socket Listener will hang or fail rather than report a clean error. This test cannot currently pass as written.

Either register a --kd_alpha arg that overrides TrainingArgs.alpha, or drop it from the test command. Fix this →


🟠 Correctness

2. make_conversation_labels can write to labels[-1] via negative indexing
dataset/targets.py:40-42:

foriinrange(n-1):
ifi<len(assistant_mask) andassistant_mask[i]:
labels[i-1] =tokens[i]

When i == 0 and assistant_mask[0] is truthy, labels[-1] = tokens[0] silently writes to the last label position (Python negative index) instead of being a no-op. In practice position 0 is usually BOS/system so the mask is 0 there, but this is a latent data-corruption bug — one mislabeled sample can inject a spurious supervised target at the final position. Guard with if i > 0: (or start the range at 1). Fix this →

3. Backend qconfig change affects existing QAT callers, not just this feature
backends/qualcomm/quantizer/qconfig.py swaps FusedMovingAvgObsFakeQuantizeFakeQuantize and MovingAverage*ObserverMinMaxObserver/PerChannelParamObserver in get_16a4w_qnn_qat_config and get_qat_per_channel_quant_config. These functions are also used by annotators/htp_rules.py and annotators/lpai_rules.py. Switching from moving-average to plain min/max observers changes calibration behavior for all existing consumers of these QAT configs. This aligns with the new strategy (observers are frozen after the calibration pass via _enable_fake_quant_freeze_observer), but it's an implicit behavior change outside the LLM example scope. Please confirm this is intended for the htp/lpai rule paths and call it out in the PR summary.


🟡 Minor / Questions

4. Operator-precedence readability in the use_qat guardllm_wrappers.py:599-603:

if (
self.mode==Mode.CALIBRATEandself.control_args.quant_recipe_suggestionoruse_qat
):

This parses as (mode==CALIBRATE and quant_recipe_suggestion) or use_qat, which is correct (use_qat already implies CALIBRATE), but the mixed and/or without parens is easy to misread. Wrap the and clause in explicit parentheses.

5. _split_train_val / _split_qat_dataset can produce an all-drop_last=True empty loader. With small --train_limit (the CI uses train_limit=1) and batch_size=1 this is fine, but n_train = max(1, ...) followed by drop_last=True means a training subset with a single sample that doesn't fill batch_size>1 yields zero batches → len(train_loader)==0total_steps==0. _log_dataset_stats guards the calib path but the post-split train subset isn't re-checked. Consider validating len(train_loader) > 0 after the split.

6. KLDivergenceLoss zero-path device but not dtypeloss.py:59 returns torch.zeros(1, device=...).squeeze() (float32) as a fallback; fine on CPU, but if the model is ever autocast/half this won't match student_logits.dtype. Low priority since device=torch.device("cpu") is hardcoded today.

7. BaseTrainer.device is hardcoded to CPU (trainer.py:45, with a # TODO: validate GPU compatibility). Worth noting in the guidance doc that QAT is CPU-only today — for larger decoders this training pass could be very slow.

8. Stale docstring in build_qat_dataloaders references --calib_train_ratio / --train_val_ratio for the "full" mode; the explicit-split branch docstring is accurate, but a reader skimming the method header sees "Full split" first. Minor.


Nits

  • llama.py:528 adds a stray blank line after --r3 default.
  • Help text typos: "lm-eluther"lm-eval / lm-evaluation-harness (appears in --qat_full_tasks and --train_tasks help).
  • train/config/config.py docstring says defaults live in train/config/config.py but the parser help (llama.py:659) points to train/config/config.py too — consistent, good; just confirm the qad.yaml comment "train/config.py" (missing config/) is updated.

Overall the design is solid and the abstractions are clean. Item #1 blocks the test from passing and #2/#3 are the substantive correctness items to resolve before merge.
· branch dev1/danny/llm_QAT_QAD

@psiddh

Copy link
Copy Markdown
Contributor

@DannyYuyang-quic can you resolve teh merge conflicts ?

…for LLMs
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
Testing:
- Add test_static_llm_qat: assert QAT PPL < PTQ PPL on smollm2_135m
@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@psiddh I've rebased the branch and addressed code review comments. Please take a look, thanks!

@psiddh
psiddh merged commit c461421 into pytorch:mainAug 13, 2026
195 of 196 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.release notes: qualcommChanges to the Qualcomm backend delegate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@DannyYuyang-quic@psiddh
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix"); } } catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

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

Merged
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD
Aug 13, 2026
Merged

Qualcomm AI Engine Direct - [LLM QAT] LLM Quant-Aware Distillation (QAD)#21036
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD

Conversation

@DannyYuyang-quic

@DannyYuyang-quicDannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

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:

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

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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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

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

@pytorch-bot

pytorch-botBot commented Jul 20, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21036

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job

As of commit 54a691b with merge base bc2833f (image):

CANCELLED JOB - The following job was cancelled. Please retry:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 20, 2026
@DannyYuyang-quic

DannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
ContributorAuthor

@psiddh Hi,
As more users experiment with lower-bit quantization, maintaining model quality with PTQ alone becomes increasingly challenging. To improve output quality under more aggressive quantization settings, we've explored and added a simple end-to-end QAT (Quantization-Aware Training) flow for LLMs.

With QAT, we can further push quantization down to W4 PCQ, and potentially even lower precisions (W2). Below is a comparison between PTQ and QAT on SmolLM2-135M-instruct using W4 per-channel quantization for all decoder layers.

Please have a look, thanks!

Experiment Results: W4 Decoder PCQ (QAT vs PTQ)

QAT (W4 PCQ decoder)

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

[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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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|>

PTQ (W4 PCQ decoder)

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?" --calib_hf_dataset "HuggingFaceTB/smol-smoltalk" --calib_hf_limit 1000

Result

[INFO 2026-07-20 15:12:57,609 llama.py:290] Device Inference Results[0]:
<|im_start|>user
What dose it mean to edit written content?<|im_end|><|im_start|>assistant
1. A sentence is usually written in the first person, i.e., "He said, she said" or "She said, she said."
2. A paragraph is usually written in the third person, i.e., "She said, he said" or "She said, she said"
3. A couple is usually written in the second person, i.e., "She said, he said" or "He said, she said"
4. A little is usually written in the third person, i.e., "She said, he said" or "She said, she said"
5. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
6. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
7. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
8. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
9. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
10. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "He said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "

@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@pytorchbot label "release notes: qualcomm"

@pytorch-botpytorch-botBot added the release notes: qualcomm Changes to the Qualcomm backend delegate label Jul 20, 2026
@psiddh

Copy link
Copy Markdown
Contributor

@claude review

@psiddh

Copy link
Copy Markdown
Contributor

@claude

claudeBot commented Aug 11, 2026

Copy link
Copy Markdown

Claude finished @psiddh's task in 5m 52s —— View job


Claude Code Review

Reviewed the full diff (git diff origin/main...HEAD, 35 files). This is a well-structured QAT/QAD feature: the trainer abstraction, injected forward_fn, and PTQ-then-train strategy composition are clean and the CI test that asserts QAT PPL < PTQ PPL is a nice regression guard. Below are the issues I found, ordered by severity.


🔴 Blocking

1. CI test passes --kd_alpha, an argument that doesn't exist
backends/qualcomm/tests/test_qnn_delegate.py:8814 appends "--kd_alpha", "0.5" to the command, but there is no --kd_alpha in _build_parser() (the KD weight is alpha, sourced only from the --train_config YAML). Since llama.py:1022 uses strict parser.parse_args(), the subprocess will exit with error: unrecognized arguments: --kd_alpha 0.5, and the test's socket Listener will hang or fail rather than report a clean error. This test cannot currently pass as written.

Either register a --kd_alpha arg that overrides TrainingArgs.alpha, or drop it from the test command. Fix this →


🟠 Correctness

2. make_conversation_labels can write to labels[-1] via negative indexing
dataset/targets.py:40-42:

foriinrange(n-1):
ifi<len(assistant_mask) andassistant_mask[i]:
labels[i-1] =tokens[i]

When i == 0 and assistant_mask[0] is truthy, labels[-1] = tokens[0] silently writes to the last label position (Python negative index) instead of being a no-op. In practice position 0 is usually BOS/system so the mask is 0 there, but this is a latent data-corruption bug — one mislabeled sample can inject a spurious supervised target at the final position. Guard with if i > 0: (or start the range at 1). Fix this →

3. Backend qconfig change affects existing QAT callers, not just this feature
backends/qualcomm/quantizer/qconfig.py swaps FusedMovingAvgObsFakeQuantizeFakeQuantize and MovingAverage*ObserverMinMaxObserver/PerChannelParamObserver in get_16a4w_qnn_qat_config and get_qat_per_channel_quant_config. These functions are also used by annotators/htp_rules.py and annotators/lpai_rules.py. Switching from moving-average to plain min/max observers changes calibration behavior for all existing consumers of these QAT configs. This aligns with the new strategy (observers are frozen after the calibration pass via _enable_fake_quant_freeze_observer), but it's an implicit behavior change outside the LLM example scope. Please confirm this is intended for the htp/lpai rule paths and call it out in the PR summary.


🟡 Minor / Questions

4. Operator-precedence readability in the use_qat guardllm_wrappers.py:599-603:

if (
self.mode==Mode.CALIBRATEandself.control_args.quant_recipe_suggestionoruse_qat
):

This parses as (mode==CALIBRATE and quant_recipe_suggestion) or use_qat, which is correct (use_qat already implies CALIBRATE), but the mixed and/or without parens is easy to misread. Wrap the and clause in explicit parentheses.

5. _split_train_val / _split_qat_dataset can produce an all-drop_last=True empty loader. With small --train_limit (the CI uses train_limit=1) and batch_size=1 this is fine, but n_train = max(1, ...) followed by drop_last=True means a training subset with a single sample that doesn't fill batch_size>1 yields zero batches → len(train_loader)==0total_steps==0. _log_dataset_stats guards the calib path but the post-split train subset isn't re-checked. Consider validating len(train_loader) > 0 after the split.

6. KLDivergenceLoss zero-path device but not dtypeloss.py:59 returns torch.zeros(1, device=...).squeeze() (float32) as a fallback; fine on CPU, but if the model is ever autocast/half this won't match student_logits.dtype. Low priority since device=torch.device("cpu") is hardcoded today.

7. BaseTrainer.device is hardcoded to CPU (trainer.py:45, with a # TODO: validate GPU compatibility). Worth noting in the guidance doc that QAT is CPU-only today — for larger decoders this training pass could be very slow.

8. Stale docstring in build_qat_dataloaders references --calib_train_ratio / --train_val_ratio for the "full" mode; the explicit-split branch docstring is accurate, but a reader skimming the method header sees "Full split" first. Minor.


Nits

  • llama.py:528 adds a stray blank line after --r3 default.
  • Help text typos: "lm-eluther"lm-eval / lm-evaluation-harness (appears in --qat_full_tasks and --train_tasks help).
  • train/config/config.py docstring says defaults live in train/config/config.py but the parser help (llama.py:659) points to train/config/config.py too — consistent, good; just confirm the qad.yaml comment "train/config.py" (missing config/) is updated.

Overall the design is solid and the abstractions are clean. Item #1 blocks the test from passing and #2/#3 are the substantive correctness items to resolve before merge.
· branch dev1/danny/llm_QAT_QAD

@psiddh

Copy link
Copy Markdown
Contributor

@DannyYuyang-quic can you resolve teh merge conflicts ?

…for LLMs
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
Testing:
- Add test_static_llm_qat: assert QAT PPL < PTQ PPL on smollm2_135m
@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@psiddh I've rebased the branch and addressed code review comments. Please take a look, thanks!

@psiddh
psiddh merged commit c461421 into pytorch:mainAug 13, 2026
195 of 196 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.release notes: qualcommChanges to the Qualcomm backend delegate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@DannyYuyang-quic@psiddh
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms"); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

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

Merged
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD
Aug 13, 2026
Merged

Qualcomm AI Engine Direct - [LLM QAT] LLM Quant-Aware Distillation (QAD)#21036
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD

Conversation

@DannyYuyang-quic

@DannyYuyang-quicDannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

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:

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

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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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

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

@pytorch-bot

pytorch-botBot commented Jul 20, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21036

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job

As of commit 54a691b with merge base bc2833f (image):

CANCELLED JOB - The following job was cancelled. Please retry:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 20, 2026
@DannyYuyang-quic

DannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
ContributorAuthor

@psiddh Hi,
As more users experiment with lower-bit quantization, maintaining model quality with PTQ alone becomes increasingly challenging. To improve output quality under more aggressive quantization settings, we've explored and added a simple end-to-end QAT (Quantization-Aware Training) flow for LLMs.

With QAT, we can further push quantization down to W4 PCQ, and potentially even lower precisions (W2). Below is a comparison between PTQ and QAT on SmolLM2-135M-instruct using W4 per-channel quantization for all decoder layers.

Please have a look, thanks!

Experiment Results: W4 Decoder PCQ (QAT vs PTQ)

QAT (W4 PCQ decoder)

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

[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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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|>

PTQ (W4 PCQ decoder)

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?" --calib_hf_dataset "HuggingFaceTB/smol-smoltalk" --calib_hf_limit 1000

Result

[INFO 2026-07-20 15:12:57,609 llama.py:290] Device Inference Results[0]:
<|im_start|>user
What dose it mean to edit written content?<|im_end|><|im_start|>assistant
1. A sentence is usually written in the first person, i.e., "He said, she said" or "She said, she said."
2. A paragraph is usually written in the third person, i.e., "She said, he said" or "She said, she said"
3. A couple is usually written in the second person, i.e., "She said, he said" or "He said, she said"
4. A little is usually written in the third person, i.e., "She said, he said" or "She said, she said"
5. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
6. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
7. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
8. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
9. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
10. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "He said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "

@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@pytorchbot label "release notes: qualcomm"

@pytorch-botpytorch-botBot added the release notes: qualcomm Changes to the Qualcomm backend delegate label Jul 20, 2026
@psiddh

Copy link
Copy Markdown
Contributor

@claude review

@psiddh

Copy link
Copy Markdown
Contributor

@claude

claudeBot commented Aug 11, 2026

Copy link
Copy Markdown

Claude finished @psiddh's task in 5m 52s —— View job


Claude Code Review

Reviewed the full diff (git diff origin/main...HEAD, 35 files). This is a well-structured QAT/QAD feature: the trainer abstraction, injected forward_fn, and PTQ-then-train strategy composition are clean and the CI test that asserts QAT PPL < PTQ PPL is a nice regression guard. Below are the issues I found, ordered by severity.


🔴 Blocking

1. CI test passes --kd_alpha, an argument that doesn't exist
backends/qualcomm/tests/test_qnn_delegate.py:8814 appends "--kd_alpha", "0.5" to the command, but there is no --kd_alpha in _build_parser() (the KD weight is alpha, sourced only from the --train_config YAML). Since llama.py:1022 uses strict parser.parse_args(), the subprocess will exit with error: unrecognized arguments: --kd_alpha 0.5, and the test's socket Listener will hang or fail rather than report a clean error. This test cannot currently pass as written.

Either register a --kd_alpha arg that overrides TrainingArgs.alpha, or drop it from the test command. Fix this →


🟠 Correctness

2. make_conversation_labels can write to labels[-1] via negative indexing
dataset/targets.py:40-42:

foriinrange(n-1):
ifi<len(assistant_mask) andassistant_mask[i]:
labels[i-1] =tokens[i]

When i == 0 and assistant_mask[0] is truthy, labels[-1] = tokens[0] silently writes to the last label position (Python negative index) instead of being a no-op. In practice position 0 is usually BOS/system so the mask is 0 there, but this is a latent data-corruption bug — one mislabeled sample can inject a spurious supervised target at the final position. Guard with if i > 0: (or start the range at 1). Fix this →

3. Backend qconfig change affects existing QAT callers, not just this feature
backends/qualcomm/quantizer/qconfig.py swaps FusedMovingAvgObsFakeQuantizeFakeQuantize and MovingAverage*ObserverMinMaxObserver/PerChannelParamObserver in get_16a4w_qnn_qat_config and get_qat_per_channel_quant_config. These functions are also used by annotators/htp_rules.py and annotators/lpai_rules.py. Switching from moving-average to plain min/max observers changes calibration behavior for all existing consumers of these QAT configs. This aligns with the new strategy (observers are frozen after the calibration pass via _enable_fake_quant_freeze_observer), but it's an implicit behavior change outside the LLM example scope. Please confirm this is intended for the htp/lpai rule paths and call it out in the PR summary.


🟡 Minor / Questions

4. Operator-precedence readability in the use_qat guardllm_wrappers.py:599-603:

if (
self.mode==Mode.CALIBRATEandself.control_args.quant_recipe_suggestionoruse_qat
):

This parses as (mode==CALIBRATE and quant_recipe_suggestion) or use_qat, which is correct (use_qat already implies CALIBRATE), but the mixed and/or without parens is easy to misread. Wrap the and clause in explicit parentheses.

5. _split_train_val / _split_qat_dataset can produce an all-drop_last=True empty loader. With small --train_limit (the CI uses train_limit=1) and batch_size=1 this is fine, but n_train = max(1, ...) followed by drop_last=True means a training subset with a single sample that doesn't fill batch_size>1 yields zero batches → len(train_loader)==0total_steps==0. _log_dataset_stats guards the calib path but the post-split train subset isn't re-checked. Consider validating len(train_loader) > 0 after the split.

6. KLDivergenceLoss zero-path device but not dtypeloss.py:59 returns torch.zeros(1, device=...).squeeze() (float32) as a fallback; fine on CPU, but if the model is ever autocast/half this won't match student_logits.dtype. Low priority since device=torch.device("cpu") is hardcoded today.

7. BaseTrainer.device is hardcoded to CPU (trainer.py:45, with a # TODO: validate GPU compatibility). Worth noting in the guidance doc that QAT is CPU-only today — for larger decoders this training pass could be very slow.

8. Stale docstring in build_qat_dataloaders references --calib_train_ratio / --train_val_ratio for the "full" mode; the explicit-split branch docstring is accurate, but a reader skimming the method header sees "Full split" first. Minor.


Nits

  • llama.py:528 adds a stray blank line after --r3 default.
  • Help text typos: "lm-eluther"lm-eval / lm-evaluation-harness (appears in --qat_full_tasks and --train_tasks help).
  • train/config/config.py docstring says defaults live in train/config/config.py but the parser help (llama.py:659) points to train/config/config.py too — consistent, good; just confirm the qad.yaml comment "train/config.py" (missing config/) is updated.

Overall the design is solid and the abstractions are clean. Item #1 blocks the test from passing and #2/#3 are the substantive correctness items to resolve before merge.
· branch dev1/danny/llm_QAT_QAD

@psiddh

Copy link
Copy Markdown
Contributor

@DannyYuyang-quic can you resolve teh merge conflicts ?

…for LLMs
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
Testing:
- Add test_static_llm_qat: assert QAT PPL < PTQ PPL on smollm2_135m
@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@psiddh I've rebased the branch and addressed code review comments. Please take a look, thanks!

@psiddh
psiddh merged commit c461421 into pytorch:mainAug 13, 2026
195 of 196 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.release notes: qualcommChanges to the Qualcomm backend delegate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@DannyYuyang-quic@psiddh
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links"); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "youtube.com"; var __re = new RegExp('^' + "youtube\\.com" + '
Skip to content

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

Merged
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD
Aug 13, 2026
Merged

Qualcomm AI Engine Direct - [LLM QAT] LLM Quant-Aware Distillation (QAD)#21036
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD

Conversation

@DannyYuyang-quic

@DannyYuyang-quicDannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

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:

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

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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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

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

@pytorch-bot

pytorch-botBot commented Jul 20, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21036

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job

As of commit 54a691b with merge base bc2833f (image):

CANCELLED JOB - The following job was cancelled. Please retry:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 20, 2026
@DannyYuyang-quic

DannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
ContributorAuthor

@psiddh Hi,
As more users experiment with lower-bit quantization, maintaining model quality with PTQ alone becomes increasingly challenging. To improve output quality under more aggressive quantization settings, we've explored and added a simple end-to-end QAT (Quantization-Aware Training) flow for LLMs.

With QAT, we can further push quantization down to W4 PCQ, and potentially even lower precisions (W2). Below is a comparison between PTQ and QAT on SmolLM2-135M-instruct using W4 per-channel quantization for all decoder layers.

Please have a look, thanks!

Experiment Results: W4 Decoder PCQ (QAT vs PTQ)

QAT (W4 PCQ decoder)

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

[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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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|>

PTQ (W4 PCQ decoder)

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?" --calib_hf_dataset "HuggingFaceTB/smol-smoltalk" --calib_hf_limit 1000

Result

[INFO 2026-07-20 15:12:57,609 llama.py:290] Device Inference Results[0]:
<|im_start|>user
What dose it mean to edit written content?<|im_end|><|im_start|>assistant
1. A sentence is usually written in the first person, i.e., "He said, she said" or "She said, she said."
2. A paragraph is usually written in the third person, i.e., "She said, he said" or "She said, she said"
3. A couple is usually written in the second person, i.e., "She said, he said" or "He said, she said"
4. A little is usually written in the third person, i.e., "She said, he said" or "She said, she said"
5. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
6. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
7. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
8. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
9. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
10. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "He said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "

@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@pytorchbot label "release notes: qualcomm"

@pytorch-botpytorch-botBot added the release notes: qualcomm Changes to the Qualcomm backend delegate label Jul 20, 2026
@psiddh

Copy link
Copy Markdown
Contributor

@claude review

@psiddh

Copy link
Copy Markdown
Contributor

@claude

claudeBot commented Aug 11, 2026

Copy link
Copy Markdown

Claude finished @psiddh's task in 5m 52s —— View job


Claude Code Review

Reviewed the full diff (git diff origin/main...HEAD, 35 files). This is a well-structured QAT/QAD feature: the trainer abstraction, injected forward_fn, and PTQ-then-train strategy composition are clean and the CI test that asserts QAT PPL < PTQ PPL is a nice regression guard. Below are the issues I found, ordered by severity.


🔴 Blocking

1. CI test passes --kd_alpha, an argument that doesn't exist
backends/qualcomm/tests/test_qnn_delegate.py:8814 appends "--kd_alpha", "0.5" to the command, but there is no --kd_alpha in _build_parser() (the KD weight is alpha, sourced only from the --train_config YAML). Since llama.py:1022 uses strict parser.parse_args(), the subprocess will exit with error: unrecognized arguments: --kd_alpha 0.5, and the test's socket Listener will hang or fail rather than report a clean error. This test cannot currently pass as written.

Either register a --kd_alpha arg that overrides TrainingArgs.alpha, or drop it from the test command. Fix this →


🟠 Correctness

2. make_conversation_labels can write to labels[-1] via negative indexing
dataset/targets.py:40-42:

foriinrange(n-1):
ifi<len(assistant_mask) andassistant_mask[i]:
labels[i-1] =tokens[i]

When i == 0 and assistant_mask[0] is truthy, labels[-1] = tokens[0] silently writes to the last label position (Python negative index) instead of being a no-op. In practice position 0 is usually BOS/system so the mask is 0 there, but this is a latent data-corruption bug — one mislabeled sample can inject a spurious supervised target at the final position. Guard with if i > 0: (or start the range at 1). Fix this →

3. Backend qconfig change affects existing QAT callers, not just this feature
backends/qualcomm/quantizer/qconfig.py swaps FusedMovingAvgObsFakeQuantizeFakeQuantize and MovingAverage*ObserverMinMaxObserver/PerChannelParamObserver in get_16a4w_qnn_qat_config and get_qat_per_channel_quant_config. These functions are also used by annotators/htp_rules.py and annotators/lpai_rules.py. Switching from moving-average to plain min/max observers changes calibration behavior for all existing consumers of these QAT configs. This aligns with the new strategy (observers are frozen after the calibration pass via _enable_fake_quant_freeze_observer), but it's an implicit behavior change outside the LLM example scope. Please confirm this is intended for the htp/lpai rule paths and call it out in the PR summary.


🟡 Minor / Questions

4. Operator-precedence readability in the use_qat guardllm_wrappers.py:599-603:

if (
self.mode==Mode.CALIBRATEandself.control_args.quant_recipe_suggestionoruse_qat
):

This parses as (mode==CALIBRATE and quant_recipe_suggestion) or use_qat, which is correct (use_qat already implies CALIBRATE), but the mixed and/or without parens is easy to misread. Wrap the and clause in explicit parentheses.

5. _split_train_val / _split_qat_dataset can produce an all-drop_last=True empty loader. With small --train_limit (the CI uses train_limit=1) and batch_size=1 this is fine, but n_train = max(1, ...) followed by drop_last=True means a training subset with a single sample that doesn't fill batch_size>1 yields zero batches → len(train_loader)==0total_steps==0. _log_dataset_stats guards the calib path but the post-split train subset isn't re-checked. Consider validating len(train_loader) > 0 after the split.

6. KLDivergenceLoss zero-path device but not dtypeloss.py:59 returns torch.zeros(1, device=...).squeeze() (float32) as a fallback; fine on CPU, but if the model is ever autocast/half this won't match student_logits.dtype. Low priority since device=torch.device("cpu") is hardcoded today.

7. BaseTrainer.device is hardcoded to CPU (trainer.py:45, with a # TODO: validate GPU compatibility). Worth noting in the guidance doc that QAT is CPU-only today — for larger decoders this training pass could be very slow.

8. Stale docstring in build_qat_dataloaders references --calib_train_ratio / --train_val_ratio for the "full" mode; the explicit-split branch docstring is accurate, but a reader skimming the method header sees "Full split" first. Minor.


Nits

  • llama.py:528 adds a stray blank line after --r3 default.
  • Help text typos: "lm-eluther"lm-eval / lm-evaluation-harness (appears in --qat_full_tasks and --train_tasks help).
  • train/config/config.py docstring says defaults live in train/config/config.py but the parser help (llama.py:659) points to train/config/config.py too — consistent, good; just confirm the qad.yaml comment "train/config.py" (missing config/) is updated.

Overall the design is solid and the abstractions are clean. Item #1 blocks the test from passing and #2/#3 are the substantive correctness items to resolve before merge.
· branch dev1/danny/llm_QAT_QAD

@psiddh

Copy link
Copy Markdown
Contributor

@DannyYuyang-quic can you resolve teh merge conflicts ?

…for LLMs
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
Testing:
- Add test_static_llm_qat: assert QAT PPL < PTQ PPL on smollm2_135m
@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@psiddh I've rebased the branch and addressed code review comments. Please take a look, thanks!

@psiddh
psiddh merged commit c461421 into pytorch:mainAug 13, 2026
195 of 196 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.release notes: qualcommChanges to the Qualcomm backend delegate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@DannyYuyang-quic@psiddh
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default"); } } catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

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

Merged
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD
Aug 13, 2026
Merged

Qualcomm AI Engine Direct - [LLM QAT] LLM Quant-Aware Distillation (QAD)#21036
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD

Conversation

@DannyYuyang-quic

@DannyYuyang-quicDannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

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:

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

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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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

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

@pytorch-bot

pytorch-botBot commented Jul 20, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21036

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job

As of commit 54a691b with merge base bc2833f (image):

CANCELLED JOB - The following job was cancelled. Please retry:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 20, 2026
@DannyYuyang-quic

DannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
ContributorAuthor

@psiddh Hi,
As more users experiment with lower-bit quantization, maintaining model quality with PTQ alone becomes increasingly challenging. To improve output quality under more aggressive quantization settings, we've explored and added a simple end-to-end QAT (Quantization-Aware Training) flow for LLMs.

With QAT, we can further push quantization down to W4 PCQ, and potentially even lower precisions (W2). Below is a comparison between PTQ and QAT on SmolLM2-135M-instruct using W4 per-channel quantization for all decoder layers.

Please have a look, thanks!

Experiment Results: W4 Decoder PCQ (QAT vs PTQ)

QAT (W4 PCQ decoder)

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

[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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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|>

PTQ (W4 PCQ decoder)

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?" --calib_hf_dataset "HuggingFaceTB/smol-smoltalk" --calib_hf_limit 1000

Result

[INFO 2026-07-20 15:12:57,609 llama.py:290] Device Inference Results[0]:
<|im_start|>user
What dose it mean to edit written content?<|im_end|><|im_start|>assistant
1. A sentence is usually written in the first person, i.e., "He said, she said" or "She said, she said."
2. A paragraph is usually written in the third person, i.e., "She said, he said" or "She said, she said"
3. A couple is usually written in the second person, i.e., "She said, he said" or "He said, she said"
4. A little is usually written in the third person, i.e., "She said, he said" or "She said, she said"
5. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
6. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
7. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
8. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
9. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
10. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "He said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "

@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@pytorchbot label "release notes: qualcomm"

@pytorch-botpytorch-botBot added the release notes: qualcomm Changes to the Qualcomm backend delegate label Jul 20, 2026
@psiddh

Copy link
Copy Markdown
Contributor

@claude review

@psiddh

Copy link
Copy Markdown
Contributor

@claude

claudeBot commented Aug 11, 2026

Copy link
Copy Markdown

Claude finished @psiddh's task in 5m 52s —— View job


Claude Code Review

Reviewed the full diff (git diff origin/main...HEAD, 35 files). This is a well-structured QAT/QAD feature: the trainer abstraction, injected forward_fn, and PTQ-then-train strategy composition are clean and the CI test that asserts QAT PPL < PTQ PPL is a nice regression guard. Below are the issues I found, ordered by severity.


🔴 Blocking

1. CI test passes --kd_alpha, an argument that doesn't exist
backends/qualcomm/tests/test_qnn_delegate.py:8814 appends "--kd_alpha", "0.5" to the command, but there is no --kd_alpha in _build_parser() (the KD weight is alpha, sourced only from the --train_config YAML). Since llama.py:1022 uses strict parser.parse_args(), the subprocess will exit with error: unrecognized arguments: --kd_alpha 0.5, and the test's socket Listener will hang or fail rather than report a clean error. This test cannot currently pass as written.

Either register a --kd_alpha arg that overrides TrainingArgs.alpha, or drop it from the test command. Fix this →


🟠 Correctness

2. make_conversation_labels can write to labels[-1] via negative indexing
dataset/targets.py:40-42:

foriinrange(n-1):
ifi<len(assistant_mask) andassistant_mask[i]:
labels[i-1] =tokens[i]

When i == 0 and assistant_mask[0] is truthy, labels[-1] = tokens[0] silently writes to the last label position (Python negative index) instead of being a no-op. In practice position 0 is usually BOS/system so the mask is 0 there, but this is a latent data-corruption bug — one mislabeled sample can inject a spurious supervised target at the final position. Guard with if i > 0: (or start the range at 1). Fix this →

3. Backend qconfig change affects existing QAT callers, not just this feature
backends/qualcomm/quantizer/qconfig.py swaps FusedMovingAvgObsFakeQuantizeFakeQuantize and MovingAverage*ObserverMinMaxObserver/PerChannelParamObserver in get_16a4w_qnn_qat_config and get_qat_per_channel_quant_config. These functions are also used by annotators/htp_rules.py and annotators/lpai_rules.py. Switching from moving-average to plain min/max observers changes calibration behavior for all existing consumers of these QAT configs. This aligns with the new strategy (observers are frozen after the calibration pass via _enable_fake_quant_freeze_observer), but it's an implicit behavior change outside the LLM example scope. Please confirm this is intended for the htp/lpai rule paths and call it out in the PR summary.


🟡 Minor / Questions

4. Operator-precedence readability in the use_qat guardllm_wrappers.py:599-603:

if (
self.mode==Mode.CALIBRATEandself.control_args.quant_recipe_suggestionoruse_qat
):

This parses as (mode==CALIBRATE and quant_recipe_suggestion) or use_qat, which is correct (use_qat already implies CALIBRATE), but the mixed and/or without parens is easy to misread. Wrap the and clause in explicit parentheses.

5. _split_train_val / _split_qat_dataset can produce an all-drop_last=True empty loader. With small --train_limit (the CI uses train_limit=1) and batch_size=1 this is fine, but n_train = max(1, ...) followed by drop_last=True means a training subset with a single sample that doesn't fill batch_size>1 yields zero batches → len(train_loader)==0total_steps==0. _log_dataset_stats guards the calib path but the post-split train subset isn't re-checked. Consider validating len(train_loader) > 0 after the split.

6. KLDivergenceLoss zero-path device but not dtypeloss.py:59 returns torch.zeros(1, device=...).squeeze() (float32) as a fallback; fine on CPU, but if the model is ever autocast/half this won't match student_logits.dtype. Low priority since device=torch.device("cpu") is hardcoded today.

7. BaseTrainer.device is hardcoded to CPU (trainer.py:45, with a # TODO: validate GPU compatibility). Worth noting in the guidance doc that QAT is CPU-only today — for larger decoders this training pass could be very slow.

8. Stale docstring in build_qat_dataloaders references --calib_train_ratio / --train_val_ratio for the "full" mode; the explicit-split branch docstring is accurate, but a reader skimming the method header sees "Full split" first. Minor.


Nits

  • llama.py:528 adds a stray blank line after --r3 default.
  • Help text typos: "lm-eluther"lm-eval / lm-evaluation-harness (appears in --qat_full_tasks and --train_tasks help).
  • train/config/config.py docstring says defaults live in train/config/config.py but the parser help (llama.py:659) points to train/config/config.py too — consistent, good; just confirm the qad.yaml comment "train/config.py" (missing config/) is updated.

Overall the design is solid and the abstractions are clean. Item #1 blocks the test from passing and #2/#3 are the substantive correctness items to resolve before merge.
· branch dev1/danny/llm_QAT_QAD

@psiddh

Copy link
Copy Markdown
Contributor

@DannyYuyang-quic can you resolve teh merge conflicts ?

…for LLMs
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
Testing:
- Add test_static_llm_qat: assert QAT PPL < PTQ PPL on smollm2_135m
@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@psiddh I've rebased the branch and addressed code review comments. Please take a look, thanks!

@psiddh
psiddh merged commit c461421 into pytorch:mainAug 13, 2026
195 of 196 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.release notes: qualcommChanges to the Qualcomm backend delegate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@DannyYuyang-quic@psiddh
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

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

Merged
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD
Aug 13, 2026
Merged

Qualcomm AI Engine Direct - [LLM QAT] LLM Quant-Aware Distillation (QAD)#21036
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD

Conversation

@DannyYuyang-quic

@DannyYuyang-quicDannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

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:

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

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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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

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

@pytorch-bot

pytorch-botBot commented Jul 20, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21036

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job

As of commit 54a691b with merge base bc2833f (image):

CANCELLED JOB - The following job was cancelled. Please retry:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 20, 2026
@DannyYuyang-quic

DannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
ContributorAuthor

@psiddh Hi,
As more users experiment with lower-bit quantization, maintaining model quality with PTQ alone becomes increasingly challenging. To improve output quality under more aggressive quantization settings, we've explored and added a simple end-to-end QAT (Quantization-Aware Training) flow for LLMs.

With QAT, we can further push quantization down to W4 PCQ, and potentially even lower precisions (W2). Below is a comparison between PTQ and QAT on SmolLM2-135M-instruct using W4 per-channel quantization for all decoder layers.

Please have a look, thanks!

Experiment Results: W4 Decoder PCQ (QAT vs PTQ)

QAT (W4 PCQ decoder)

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

[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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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|>

PTQ (W4 PCQ decoder)

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?" --calib_hf_dataset "HuggingFaceTB/smol-smoltalk" --calib_hf_limit 1000

Result

[INFO 2026-07-20 15:12:57,609 llama.py:290] Device Inference Results[0]:
<|im_start|>user
What dose it mean to edit written content?<|im_end|><|im_start|>assistant
1. A sentence is usually written in the first person, i.e., "He said, she said" or "She said, she said."
2. A paragraph is usually written in the third person, i.e., "She said, he said" or "She said, she said"
3. A couple is usually written in the second person, i.e., "She said, he said" or "He said, she said"
4. A little is usually written in the third person, i.e., "She said, he said" or "She said, she said"
5. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
6. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
7. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
8. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
9. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
10. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "He said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "

@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@pytorchbot label "release notes: qualcomm"

@pytorch-botpytorch-botBot added the release notes: qualcomm Changes to the Qualcomm backend delegate label Jul 20, 2026
@psiddh

Copy link
Copy Markdown
Contributor

@claude review

@psiddh

Copy link
Copy Markdown
Contributor

@claude

claudeBot commented Aug 11, 2026

Copy link
Copy Markdown

Claude finished @psiddh's task in 5m 52s —— View job


Claude Code Review

Reviewed the full diff (git diff origin/main...HEAD, 35 files). This is a well-structured QAT/QAD feature: the trainer abstraction, injected forward_fn, and PTQ-then-train strategy composition are clean and the CI test that asserts QAT PPL < PTQ PPL is a nice regression guard. Below are the issues I found, ordered by severity.


🔴 Blocking

1. CI test passes --kd_alpha, an argument that doesn't exist
backends/qualcomm/tests/test_qnn_delegate.py:8814 appends "--kd_alpha", "0.5" to the command, but there is no --kd_alpha in _build_parser() (the KD weight is alpha, sourced only from the --train_config YAML). Since llama.py:1022 uses strict parser.parse_args(), the subprocess will exit with error: unrecognized arguments: --kd_alpha 0.5, and the test's socket Listener will hang or fail rather than report a clean error. This test cannot currently pass as written.

Either register a --kd_alpha arg that overrides TrainingArgs.alpha, or drop it from the test command. Fix this →


🟠 Correctness

2. make_conversation_labels can write to labels[-1] via negative indexing
dataset/targets.py:40-42:

foriinrange(n-1):
ifi<len(assistant_mask) andassistant_mask[i]:
labels[i-1] =tokens[i]

When i == 0 and assistant_mask[0] is truthy, labels[-1] = tokens[0] silently writes to the last label position (Python negative index) instead of being a no-op. In practice position 0 is usually BOS/system so the mask is 0 there, but this is a latent data-corruption bug — one mislabeled sample can inject a spurious supervised target at the final position. Guard with if i > 0: (or start the range at 1). Fix this →

3. Backend qconfig change affects existing QAT callers, not just this feature
backends/qualcomm/quantizer/qconfig.py swaps FusedMovingAvgObsFakeQuantizeFakeQuantize and MovingAverage*ObserverMinMaxObserver/PerChannelParamObserver in get_16a4w_qnn_qat_config and get_qat_per_channel_quant_config. These functions are also used by annotators/htp_rules.py and annotators/lpai_rules.py. Switching from moving-average to plain min/max observers changes calibration behavior for all existing consumers of these QAT configs. This aligns with the new strategy (observers are frozen after the calibration pass via _enable_fake_quant_freeze_observer), but it's an implicit behavior change outside the LLM example scope. Please confirm this is intended for the htp/lpai rule paths and call it out in the PR summary.


🟡 Minor / Questions

4. Operator-precedence readability in the use_qat guardllm_wrappers.py:599-603:

if (
self.mode==Mode.CALIBRATEandself.control_args.quant_recipe_suggestionoruse_qat
):

This parses as (mode==CALIBRATE and quant_recipe_suggestion) or use_qat, which is correct (use_qat already implies CALIBRATE), but the mixed and/or without parens is easy to misread. Wrap the and clause in explicit parentheses.

5. _split_train_val / _split_qat_dataset can produce an all-drop_last=True empty loader. With small --train_limit (the CI uses train_limit=1) and batch_size=1 this is fine, but n_train = max(1, ...) followed by drop_last=True means a training subset with a single sample that doesn't fill batch_size>1 yields zero batches → len(train_loader)==0total_steps==0. _log_dataset_stats guards the calib path but the post-split train subset isn't re-checked. Consider validating len(train_loader) > 0 after the split.

6. KLDivergenceLoss zero-path device but not dtypeloss.py:59 returns torch.zeros(1, device=...).squeeze() (float32) as a fallback; fine on CPU, but if the model is ever autocast/half this won't match student_logits.dtype. Low priority since device=torch.device("cpu") is hardcoded today.

7. BaseTrainer.device is hardcoded to CPU (trainer.py:45, with a # TODO: validate GPU compatibility). Worth noting in the guidance doc that QAT is CPU-only today — for larger decoders this training pass could be very slow.

8. Stale docstring in build_qat_dataloaders references --calib_train_ratio / --train_val_ratio for the "full" mode; the explicit-split branch docstring is accurate, but a reader skimming the method header sees "Full split" first. Minor.


Nits

  • llama.py:528 adds a stray blank line after --r3 default.
  • Help text typos: "lm-eluther"lm-eval / lm-evaluation-harness (appears in --qat_full_tasks and --train_tasks help).
  • train/config/config.py docstring says defaults live in train/config/config.py but the parser help (llama.py:659) points to train/config/config.py too — consistent, good; just confirm the qad.yaml comment "train/config.py" (missing config/) is updated.

Overall the design is solid and the abstractions are clean. Item #1 blocks the test from passing and #2/#3 are the substantive correctness items to resolve before merge.
· branch dev1/danny/llm_QAT_QAD

@psiddh

Copy link
Copy Markdown
Contributor

@DannyYuyang-quic can you resolve teh merge conflicts ?

…for LLMs
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
Testing:
- Add test_static_llm_qat: assert QAT PPL < PTQ PPL on smollm2_135m
@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@psiddh I've rebased the branch and addressed code review comments. Please take a look, thanks!

@psiddh
psiddh merged commit c461421 into pytorch:mainAug 13, 2026
195 of 196 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.release notes: qualcommChanges to the Qualcomm backend delegate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@DannyYuyang-quic@psiddh
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode"); } } catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); } })(); })();
Skip to content

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

Merged
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD
Aug 13, 2026
Merged

Qualcomm AI Engine Direct - [LLM QAT] LLM Quant-Aware Distillation (QAD)#21036
psiddh merged 2 commits into
pytorch:mainfrom
CodeLinaro:dev1/danny/llm_QAT_QAD

Conversation

@DannyYuyang-quic

@DannyYuyang-quicDannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

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:

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

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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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

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

@pytorch-bot

pytorch-botBot commented Jul 20, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21036

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job

As of commit 54a691b with merge base bc2833f (image):

CANCELLED JOB - The following job was cancelled. Please retry:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 20, 2026
@DannyYuyang-quic

DannyYuyang-quic commented Jul 20, 2026

Copy link
Copy Markdown
ContributorAuthor

@psiddh Hi,
As more users experiment with lower-bit quantization, maintaining model quality with PTQ alone becomes increasingly challenging. To improve output quality under more aggressive quantization settings, we've explored and added a simple end-to-end QAT (Quantization-Aware Training) flow for LLMs.

With QAT, we can further push quantization down to W4 PCQ, and potentially even lower precisions (W2). Below is a comparison between PTQ and QAT on SmolLM2-135M-instruct using W4 per-channel quantization for all decoder layers.

Please have a look, thanks!

Experiment Results: W4 Decoder PCQ (QAT vs PTQ)

QAT (W4 PCQ decoder)

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

[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 forgrammar**: Grammar is another crucial aspect of writing. When you edit your writing, you check for errorsin grammar, syntax, and punctuation, which can make your writing more polished and error-free.
3. **Editing forstyle**: 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 engagingin the way you want it to be.
4. **Editing forimpact**: 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 itin a way that resonates with your readers.
5. **Editing forconsistency**: When you edit your writing, you ensure that it is consistentin 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|>

PTQ (W4 PCQ decoder)

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?" --calib_hf_dataset "HuggingFaceTB/smol-smoltalk" --calib_hf_limit 1000

Result

[INFO 2026-07-20 15:12:57,609 llama.py:290] Device Inference Results[0]:
<|im_start|>user
What dose it mean to edit written content?<|im_end|><|im_start|>assistant
1. A sentence is usually written in the first person, i.e., "He said, she said" or "She said, she said."
2. A paragraph is usually written in the third person, i.e., "She said, he said" or "She said, she said"
3. A couple is usually written in the second person, i.e., "She said, he said" or "He said, she said"
4. A little is usually written in the third person, i.e., "She said, he said" or "She said, she said"
5. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
6. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
7. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
8. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
9. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
10. A little is usually written in the second person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "He said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 1st person, i.e., "She said, she said" or "She said, she said"
Please note that the above examples are from the 2nd person, i.e., "

@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@pytorchbot label "release notes: qualcomm"

@pytorch-botpytorch-botBot added the release notes: qualcomm Changes to the Qualcomm backend delegate label Jul 20, 2026
@psiddh

Copy link
Copy Markdown
Contributor

@claude review

@psiddh

Copy link
Copy Markdown
Contributor

@claude

claudeBot commented Aug 11, 2026

Copy link
Copy Markdown

Claude finished @psiddh's task in 5m 52s —— View job


Claude Code Review

Reviewed the full diff (git diff origin/main...HEAD, 35 files). This is a well-structured QAT/QAD feature: the trainer abstraction, injected forward_fn, and PTQ-then-train strategy composition are clean and the CI test that asserts QAT PPL < PTQ PPL is a nice regression guard. Below are the issues I found, ordered by severity.


🔴 Blocking

1. CI test passes --kd_alpha, an argument that doesn't exist
backends/qualcomm/tests/test_qnn_delegate.py:8814 appends "--kd_alpha", "0.5" to the command, but there is no --kd_alpha in _build_parser() (the KD weight is alpha, sourced only from the --train_config YAML). Since llama.py:1022 uses strict parser.parse_args(), the subprocess will exit with error: unrecognized arguments: --kd_alpha 0.5, and the test's socket Listener will hang or fail rather than report a clean error. This test cannot currently pass as written.

Either register a --kd_alpha arg that overrides TrainingArgs.alpha, or drop it from the test command. Fix this →


🟠 Correctness

2. make_conversation_labels can write to labels[-1] via negative indexing
dataset/targets.py:40-42:

foriinrange(n-1):
ifi<len(assistant_mask) andassistant_mask[i]:
labels[i-1] =tokens[i]

When i == 0 and assistant_mask[0] is truthy, labels[-1] = tokens[0] silently writes to the last label position (Python negative index) instead of being a no-op. In practice position 0 is usually BOS/system so the mask is 0 there, but this is a latent data-corruption bug — one mislabeled sample can inject a spurious supervised target at the final position. Guard with if i > 0: (or start the range at 1). Fix this →

3. Backend qconfig change affects existing QAT callers, not just this feature
backends/qualcomm/quantizer/qconfig.py swaps FusedMovingAvgObsFakeQuantizeFakeQuantize and MovingAverage*ObserverMinMaxObserver/PerChannelParamObserver in get_16a4w_qnn_qat_config and get_qat_per_channel_quant_config. These functions are also used by annotators/htp_rules.py and annotators/lpai_rules.py. Switching from moving-average to plain min/max observers changes calibration behavior for all existing consumers of these QAT configs. This aligns with the new strategy (observers are frozen after the calibration pass via _enable_fake_quant_freeze_observer), but it's an implicit behavior change outside the LLM example scope. Please confirm this is intended for the htp/lpai rule paths and call it out in the PR summary.


🟡 Minor / Questions

4. Operator-precedence readability in the use_qat guardllm_wrappers.py:599-603:

if (
self.mode==Mode.CALIBRATEandself.control_args.quant_recipe_suggestionoruse_qat
):

This parses as (mode==CALIBRATE and quant_recipe_suggestion) or use_qat, which is correct (use_qat already implies CALIBRATE), but the mixed and/or without parens is easy to misread. Wrap the and clause in explicit parentheses.

5. _split_train_val / _split_qat_dataset can produce an all-drop_last=True empty loader. With small --train_limit (the CI uses train_limit=1) and batch_size=1 this is fine, but n_train = max(1, ...) followed by drop_last=True means a training subset with a single sample that doesn't fill batch_size>1 yields zero batches → len(train_loader)==0total_steps==0. _log_dataset_stats guards the calib path but the post-split train subset isn't re-checked. Consider validating len(train_loader) > 0 after the split.

6. KLDivergenceLoss zero-path device but not dtypeloss.py:59 returns torch.zeros(1, device=...).squeeze() (float32) as a fallback; fine on CPU, but if the model is ever autocast/half this won't match student_logits.dtype. Low priority since device=torch.device("cpu") is hardcoded today.

7. BaseTrainer.device is hardcoded to CPU (trainer.py:45, with a # TODO: validate GPU compatibility). Worth noting in the guidance doc that QAT is CPU-only today — for larger decoders this training pass could be very slow.

8. Stale docstring in build_qat_dataloaders references --calib_train_ratio / --train_val_ratio for the "full" mode; the explicit-split branch docstring is accurate, but a reader skimming the method header sees "Full split" first. Minor.


Nits

  • llama.py:528 adds a stray blank line after --r3 default.
  • Help text typos: "lm-eluther"lm-eval / lm-evaluation-harness (appears in --qat_full_tasks and --train_tasks help).
  • train/config/config.py docstring says defaults live in train/config/config.py but the parser help (llama.py:659) points to train/config/config.py too — consistent, good; just confirm the qad.yaml comment "train/config.py" (missing config/) is updated.

Overall the design is solid and the abstractions are clean. Item #1 blocks the test from passing and #2/#3 are the substantive correctness items to resolve before merge.
· branch dev1/danny/llm_QAT_QAD

@psiddh

Copy link
Copy Markdown
Contributor

@DannyYuyang-quic can you resolve teh merge conflicts ?

…for LLMs
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
Testing:
- Add test_static_llm_qat: assert QAT PPL < PTQ PPL on smollm2_135m
@DannyYuyang-quic

Copy link
Copy Markdown
ContributorAuthor

@psiddh I've rebased the branch and addressed code review comments. Please take a look, thanks!

@psiddh
psiddh merged commit c461421 into pytorch:mainAug 13, 2026
195 of 196 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.release notes: qualcommChanges to the Qualcomm backend delegate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@DannyYuyang-quic@psiddh