diff --git a/backends/qualcomm/_passes/seq_mse.py b/backends/qualcomm/_passes/seq_mse.py index 7ec399699dd..f2b5c18b839 100644 --- a/backends/qualcomm/_passes/seq_mse.py +++ b/backends/qualcomm/_passes/seq_mse.py @@ -160,6 +160,9 @@ def _insert_seq_mse( # extract observer weight_node_obs = node.args[1] observer = getattr(graph_module, weight_node_obs.name) + # For fake quant case + if hasattr(observer, "activation_post_process"): + observer = observer.activation_post_process # extract parameters weight_node = weight_node_obs.args[0] weight_tensor = graph_module.get_parameter(weight_node.target).detach() diff --git a/backends/qualcomm/quantizer/qconfig.py b/backends/qualcomm/quantizer/qconfig.py index b75661390ca..3f523eb7d3f 100644 --- a/backends/qualcomm/quantizer/qconfig.py +++ b/backends/qualcomm/quantizer/qconfig.py @@ -896,13 +896,13 @@ def get_16a4w_qnn_qat_config( observer_or_fake_quant_ctr=act_fake_quant_ctr, ) - weight_fake_quant_ctr = FusedMovingAvgObsFakeQuantize.with_args( + weight_fake_quant_ctr = FakeQuantize.with_args( dtype=torch.int8, quant_min=-7, quant_max=7, qscheme=torch.per_tensor_symmetric, ch_axis=0, - observer=MovingAverageMinMaxObserver.with_args(**extra_args), + observer=MinMaxObserver.with_args(**extra_args), ) weight_quantization_spec = QuantizationSpec( dtype=torch.int8, @@ -918,7 +918,7 @@ def get_16a4w_qnn_qat_config( quant_min=torch.iinfo(torch.int32).min, quant_max=torch.iinfo(torch.int32).max, qscheme=torch.per_tensor_symmetric, - observer=MovingAverageMinMaxObserver.with_args(**extra_args), + observer=MinMaxObserver.with_args(**extra_args), ) bias_quantization_spec = QuantizationSpec( dtype=torch.int32, @@ -998,7 +998,7 @@ def get_qat_per_channel_quant_config( observer_or_fake_quant_ctr=act_fake_quant_ctr, ) - weight_fake_quant_ctr = FusedMovingAvgObsFakeQuantize.with_args( + weight_fake_quant_ctr = FakeQuantize.with_args( dtype=torch.int8 if weight_dtype == torch.int4 else weight_dtype, quant_min=( -7 if weight_dtype == torch.int4 else torch.iinfo(weight_dtype).min + 1 @@ -1006,7 +1006,7 @@ def get_qat_per_channel_quant_config( quant_max=7 if weight_dtype == torch.int4 else torch.iinfo(weight_dtype).max, qscheme=torch.per_channel_symmetric, ch_axis=ch_axis, - observer=MovingAveragePerChannelMinMaxObserver.with_args(**extra_args), + observer=PerChannelParamObserver.with_args(**extra_args), ) weight_quantization_spec = QuantizationSpec( dtype=torch.int8 if weight_dtype == torch.int4 else weight_dtype, diff --git a/backends/qualcomm/quantizer/quantizer.py b/backends/qualcomm/quantizer/quantizer.py index 2050aba7696..6eac3fe7e79 100644 --- a/backends/qualcomm/quantizer/quantizer.py +++ b/backends/qualcomm/quantizer/quantizer.py @@ -207,6 +207,15 @@ class QuantDtype(IntEnum): partial(get_qat_per_channel_quant_config), None, ), + (QuantDtype.use_16a8w, True): ( + get_16a8w_qnn_qat_config, + partial( + get_qat_per_channel_quant_config, + act_dtype=torch.uint16, + weight_dtype=torch.int8, + ), + None, + ), } diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index 98c5839c515..47701b36a36 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -8767,6 +8767,80 @@ def test_hf_causal_lm(self): f"Expected Output: '{golden_start_with}' Actual Output: '{model_out}'", ) + def test_static_llm_qat(self): + if not self.required_envs(): + self.skipTest("missing required envs") + if self.compile_only: + self.skipTest("tasks_eval requires on-device inference") + + def run_eval( + calib_limit: int, train_limit: int, extra_args: List[str] = None + ) -> float: + prompt = "I would like to learn python, could you teach me with a simple example?" + cmds = [ + "python", + f"{self.executorch_root}/examples/qualcomm/oss_scripts/llama/llama.py", + "--artifact", + self.artifact_dir, + "--build_folder", + self.build_folder, + "--prompt", + prompt, + "--temperature", + "0", + "--decoder_model", + "smollm2_135m", + "--model_mode", + "kv", + "--max_seq_len", + "1024", + "--max_context_len", + "1024", + "--eval_methods", + "tasks_eval", + "--eval_tasks", + "wikitext", + "--eval_limit", + "1", + "--qat", + "--calib_tasks", + "wikitext", + "--calib_limit", + str(calib_limit), + "--train_tasks", + "wikitext", + "--train_limit", + str(train_limit), + "--kd_alpha", + "0.5", + ] + if extra_args: + cmds.extend(extra_args) + self.add_default_cmds(cmds) + + p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL) + with Listener((self.ip, self.port)) as listener: + conn = listener.accept() + p.communicate() + msg = json.loads(conn.recv()) + if "Error" in msg: + self.fail( + f"smollm2_135m QAT (limit={train_limit}) failed: {msg['Error']}" + ) + return msg["wiki_ppl"] + + ptq_ppl = run_eval( + calib_limit=1, train_limit=1, extra_args=["--freeze_all_params"] + ) + qat_ppl = run_eval(calib_limit=1, train_limit=1) + logging.info(f"QAT PPL={qat_ppl:.2f}") + logging.info(f"PTQ PPL={ptq_ppl:.2f}") + self.assertLess( + qat_ppl, + ptq_ppl, + f"Expected QAT PPL ({qat_ppl:.2f}) < PTQ PPL({ptq_ppl:.2f})", + ) + class TestExampleMultimodalityScript(TestQNN): diff --git a/examples/qualcomm/oss_scripts/llama/README.md b/examples/qualcomm/oss_scripts/llama/README.md index 4522582f2f6..febf154bb92 100644 --- a/examples/qualcomm/oss_scripts/llama/README.md +++ b/examples/qualcomm/oss_scripts/llama/README.md @@ -577,8 +577,6 @@ Calibration data is required for compilation. There are two ways to supply it: For LLMs, provide at least one of the two; for multimodal models, `--calib_samples` is mandatory. -Calibration and runtime evaluation use separate flag sets and can target different tasks or limits as needed: - | Purpose | Flags | |---|---| | Calibration data (lm_eval tasks) | `--calib_tasks`, `--calib_limit`, `--calib_num_fewshot` | @@ -610,19 +608,7 @@ Ready-to-use examples for each model type are provided under `assets/samples/`: | ALM (audio) | [assets/samples/audio.json](assets/samples/audio.json) | | VLM (vision) | [assets/samples/vision.json](assets/samples/vision.json) | -#### Quantization Guidance - -To automatically identify sensitive layers and generate a mixed-precision recipe suggestion, add the `--quant_recipe_suggestion` flag. During calibration, the analyzer compares FP32 and QDQ intermediate outputs layer-by-layer using SQNR, then writes two files to the working directory: - -- `{model_name}_quantization_error.csv` — per-group SQNR statistics sorted by sensitivity (most sensitive first) -- `{model_name}_suggest_recipe.py` — ready-to-use `StaticLLMQuantRecipe` subclasses optimized to apply higher-precision quantization to the most sensitive groups. - -Example: -```bash -python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} --prompt "I would like to learn python, could you teach me with a simple example?" --temperature 0 --model_mode kv --max_seq_len 1024 --decoder_model qwen3-1_7b --calib_tasks wikitext --calib_limit 1 --quant_recipe_suggestion --compile_only -``` - -After the run, pick one of the generated classes from `qwen3-1_7b_suggest_recipe.py` as your new recipe. For a full walkthrough, see [quantization_guidance.md](quantization_guidance.md). +For the full quantization workflow: PTQ (Post Training Quantization), QAT(Quantization Aware Training / Distillation), Mixed-precision, Quantization error analysis, see [quantization_guidance.md](quantization_guidance.md). #### Use attention sink for multi-turn conversations Attention sink is a way to evict cache when maximum context length be reached. diff --git a/examples/qualcomm/oss_scripts/llama/__init__.py b/examples/qualcomm/oss_scripts/llama/__init__.py index 2e0b2278909..d88012c4ff0 100644 --- a/examples/qualcomm/oss_scripts/llama/__init__.py +++ b/examples/qualcomm/oss_scripts/llama/__init__.py @@ -83,6 +83,7 @@ Qwen2_5_1_5BQuantRecipe, Qwen3_0_6BQuantRecipe, Qwen3_1_7BQuantRecipe, + Smollm2QATQuantRecipe, Smollm2QuantRecipe, Smollm3QuantRecipe, SmolVLMQuantRecipe, @@ -125,6 +126,8 @@ class LLMModelConfig(ABC): r2: Enable SpinQuant R2 quantization optimization. r3: Enable SpinQuant R3 quantization optimization. quant_recipe: Quantization recipe to use when setting quant configs for the model. + qat_recipe: Optional QAT-specific quantization recipe. Used only when --qat is set; + if None, falls back to quant_recipe. """ repo_id: str @@ -141,6 +144,7 @@ class LLMModelConfig(ABC): r2: bool r3: bool quant_recipe: StaticLLMQuantRecipe + qat_recipe: Optional[Type[StaticLLMQuantRecipe]] = None def __str__(self): # noqa: C901 """ @@ -541,6 +545,7 @@ class Smollm2_135M(LLMModelConfig): r2 = False r3 = True quant_recipe = Smollm2QuantRecipe + qat_recipe = Smollm2QATQuantRecipe @register_llm_model("smollm3-3b") diff --git a/examples/qualcomm/oss_scripts/llama/assets/quantization/quantization_analysis_workflow.png b/examples/qualcomm/oss_scripts/llama/assets/quantization/quantization_analysis_workflow.png new file mode 100644 index 00000000000..ade101eeee7 Binary files /dev/null and b/examples/qualcomm/oss_scripts/llama/assets/quantization/quantization_analysis_workflow.png differ diff --git a/examples/qualcomm/oss_scripts/llama/assets/quantization/quantization_awre_training.png b/examples/qualcomm/oss_scripts/llama/assets/quantization/quantization_awre_training.png new file mode 100644 index 00000000000..3066d346b8b Binary files /dev/null and b/examples/qualcomm/oss_scripts/llama/assets/quantization/quantization_awre_training.png differ diff --git a/examples/qualcomm/oss_scripts/llama/assets/quantization/quantization_workflow.png b/examples/qualcomm/oss_scripts/llama/assets/quantization/quantization_workflow.png new file mode 100644 index 00000000000..68c9c3b7d8f Binary files /dev/null and b/examples/qualcomm/oss_scripts/llama/assets/quantization/quantization_workflow.png differ diff --git a/examples/qualcomm/oss_scripts/llama/dataset/__init__.py b/examples/qualcomm/oss_scripts/llama/dataset/__init__.py index 08b718698aa..e25e9887e6a 100644 --- a/examples/qualcomm/oss_scripts/llama/dataset/__init__.py +++ b/examples/qualcomm/oss_scripts/llama/dataset/__init__.py @@ -11,8 +11,12 @@ ) from executorch.examples.qualcomm.oss_scripts.llama.dataset.collators import ( LLMCalibCollator, + LLMTrainingCollator, ) from executorch.examples.qualcomm.oss_scripts.llama.dataset.config import DataConfig +from executorch.examples.qualcomm.oss_scripts.llama.dataset.constants import ( + LABEL_IGNORE_INDEX, +) from executorch.examples.qualcomm.oss_scripts.llama.dataset.datasets import ( LLMDataset, ModalityEncoderDataset, @@ -21,28 +25,40 @@ collect_lm_eval_tokens, load_audio_file, load_conversation_samples, + load_hf_chat_dataset, ) from executorch.examples.qualcomm.oss_scripts.llama.dataset.preprocessors import ( ModalityPreprocessor, preprocess_encoder_inputs, ) from executorch.examples.qualcomm.oss_scripts.llama.dataset.schema import MessageSample +from executorch.examples.qualcomm.oss_scripts.llama.dataset.targets import ( + make_causal_labels, + make_conversation_labels, +) __all__ = [ # config "DataConfig", # schema "MessageSample", + # constants + "LABEL_IGNORE_INDEX", # loaders "collect_lm_eval_tokens", "load_audio_file", "load_conversation_samples", + "load_hf_chat_dataset", + # targets (label generation — low-level functions) + "make_causal_labels", + "make_conversation_labels", # datasets "LLMDataset", "ModalityEncoderDataset", # collators "ModalityEncoderCollator", "LLMCalibCollator", + "LLMTrainingCollator", # builders "DatasetBuilder", "DecoderDatasetBuilder", diff --git a/examples/qualcomm/oss_scripts/llama/dataset/builders.py b/examples/qualcomm/oss_scripts/llama/dataset/builders.py index d31301e8303..1b53a67cadc 100644 --- a/examples/qualcomm/oss_scripts/llama/dataset/builders.py +++ b/examples/qualcomm/oss_scripts/llama/dataset/builders.py @@ -11,6 +11,7 @@ from executorch.examples.qualcomm.oss_scripts.llama import LLMModelConfig from executorch.examples.qualcomm.oss_scripts.llama.dataset.collators import ( LLMCalibCollator, + LLMTrainingCollator, ModalityEncoderCollator, ) from executorch.examples.qualcomm.oss_scripts.llama.dataset.config import DataConfig @@ -22,6 +23,7 @@ collect_lm_eval_tokens, load_audio_file, load_conversation_samples, + load_hf_chat_dataset, ) from executorch.examples.qualcomm.oss_scripts.llama.dataset.preprocessors import ( ModalityPreprocessor, @@ -41,7 +43,7 @@ ) from executorch.examples.qualcomm.oss_scripts.llama.masking_utils import AttentionMask from executorch.examples.qualcomm.oss_scripts.llama.tokenizer import TokenizerWrapper -from torch.utils.data import ConcatDataset, DataLoader, Dataset +from torch.utils.data import ConcatDataset, DataLoader, Dataset, Subset from transformers.image_utils import load_image logger = logging.getLogger(__name__) @@ -110,6 +112,26 @@ def from_lm_eval( ) return LLMDataset(sequences) + def from_hf_source( + self, + dataset_name: str, + num_samples: int, + ) -> LLMDataset: + """Build an LLMDataset from a HuggingFace dataset. + + Populates assistant_masks so training computes loss only on assistant-response + tokens (assistant-only loss), ignoring user and system turns. Plain-text corpora + with no assistant turns yield all-zero masks, which fall back to full causal + next-token labels at collation time. + """ + sequences, assistant_masks = load_hf_chat_dataset( + dataset_name, + self._tokenizer_wrapper, + self._max_context_len, + num_samples=num_samples, + ) + return LLMDataset(sequences, assistant_masks=assistant_masks) + def from_tokens(self, token_ids: List[int]) -> LLMDataset: """Wrap a single pre-tokenized sequence in an LLMDataset.""" return LLMDataset([token_ids]) @@ -174,7 +196,7 @@ def from_message_samples( class DatasetBuilder: - """Orchestrates DecoderDatasetBuilder and EncoderDatasetBuilder for calibration.""" + """Orchestrates DecoderDatasetBuilder and EncoderDatasetBuilder for calibration or training.""" def __init__( self, @@ -211,6 +233,12 @@ def __init__( self._data_config.max_context_len, self._data_config.token_dtype, ) + # Collator for QAT training: pads tokens, builds masks, and labels. + self._train_collator = LLMTrainingCollator( + self._attn_mask, + self._data_config.max_context_len, + self._data_config.token_dtype, + ) def build_calib_dataloaders(self) -> Dict[str, Optional[DataLoader]]: """Calibration DataLoaders for all modalities; all modality keys always present.""" @@ -251,13 +279,21 @@ def build_calib_dataloaders(self) -> Dict[str, Optional[DataLoader]]: if cfg.calib_tasks is not None else None ), + # build from huggingface source + ( + self._text_decoder_dataset_builder.from_hf_source( + cfg.calib_hf_dataset, cfg.calib_hf_limit + ) + if cfg.calib_hf_dataset is not None + else None + ), ], ) ) if not decoder_datasets: raise ValueError( "No calibration data specified. Provide at least one of: " - "--calib_tasks, --calib_samples" + "--calib_tasks, --calib_samples, --calib_hf_dataset." ) datasets[TEXT_DECODER] = ( @@ -283,6 +319,167 @@ def build_calib_dataloaders(self) -> Dict[str, Optional[DataLoader]]: for modality, dataset in datasets.items() } + def build_qat_dataloaders(self) -> Tuple[ + Dict[str, Optional[DataLoader]], + Dict[str, Optional[DataLoader]], + Dict[str, Optional[DataLoader]], + ]: + """Build calib/train/val DataLoaders for QAT. + + Full split: --qat_full_tasks / --qat_full_hf_dataset provide a single + pool that is automatically split into calib/train/val by --calib_train_ratio / --train_val_ratio. + + Explicit split: --calib_tasks provides the calibration set and + --train_tasks / --train_hf_dataset provide the training set. An optional + val split is cut from the training pool via --train_val_ratio. + + Returns three modality-keyed dicts: (calib_loaders, train_loaders, val_loaders). + Only TEXT_DECODER carries non-None loaders. + """ + if self._is_multimodal: + raise NotImplementedError( + "QAT is not supported for multimodal models yet. " + "build_qat_dataloaders operates on the text decoder only." + ) + + cfg = self._data_config + + if cfg.qat_mode == "full": + full_datasets: List[Dataset] = list( + filter( + None, + [ + # build from lm tasks + ( + self._text_decoder_dataset_builder.from_lm_eval( + cfg.qat_full_tasks, cfg.qat_full_limit + ) + if cfg.qat_full_tasks is not None + else None + ), + # build from huggingface source + ( + self._text_decoder_dataset_builder.from_hf_source( + cfg.qat_full_hf_dataset, + cfg.qat_full_hf_limit, + ) + if cfg.qat_full_hf_dataset is not None + else None + ), + ], + ) + ) + full_dataset = ( + ConcatDataset(full_datasets) + if len(full_datasets) > 1 + else full_datasets[0] + ) + self._log_dataset_stats( + {TEXT_DECODER: full_dataset}, cfg.batch_size, phase="qat-full" + ) + calib_loader, train_loader, val_loader = self._split_qat_dataset( + full_dataset, + calib_train_ratio=cfg.calib_train_ratio, + train_val_ratio=cfg.train_val_ratio, + batch_size=cfg.batch_size, + seed=cfg.seed, + calib_collator=self._calib_collator, + train_collator=self._train_collator, + ) + else: + # Explicit split calibration dataset and training dataset + calib_datasets: List[Dataset] = list( + filter( + None, + [ + # build from samples + ( + self._text_decoder_dataset_builder.from_conversation( + load_conversation_samples(cfg.calib_samples) + ) + if cfg.calib_samples + else None + ), + # build from lm tasks + ( + self._text_decoder_dataset_builder.from_lm_eval( + cfg.calib_tasks, cfg.calib_limit + ) + if cfg.calib_tasks is not None + else None + ), + # build from huggingface source + ( + self._text_decoder_dataset_builder.from_hf_source( + cfg.calib_hf_dataset, cfg.calib_hf_limit + ) + if cfg.calib_hf_dataset is not None + else None + ), + ], + ) + ) + train_datasets: List[Dataset] = list( + filter( + None, + [ + # build from lm tasks + ( + self._text_decoder_dataset_builder.from_lm_eval( + cfg.train_tasks, cfg.train_limit + ) + if cfg.train_tasks is not None + else None + ), + # build from huggingface source + ( + self._text_decoder_dataset_builder.from_hf_source( + cfg.train_hf_dataset, + cfg.train_hf_limit, + ) + if cfg.train_hf_dataset is not None + else None + ), + ], + ) + ) + calib_full = ( + ConcatDataset(calib_datasets) + if len(calib_datasets) > 1 + else calib_datasets[0] + ) + train_full = ( + ConcatDataset(train_datasets) + if len(train_datasets) > 1 + else train_datasets[0] + ) + self._log_dataset_stats( + {TEXT_DECODER: calib_full}, cfg.batch_size, phase="qat-calibration" + ) + self._log_dataset_stats( + {TEXT_DECODER: train_full}, cfg.batch_size, phase="qat-training" + ) + calib_loader = DataLoader( + calib_full, + batch_size=cfg.batch_size, + shuffle=False, + drop_last=cfg.batch_size > 1, + collate_fn=self._calib_collator, + ) + train_loader, val_loader = self._split_train_val( + train_full, + train_val_ratio=cfg.train_val_ratio, + batch_size=cfg.batch_size, + seed=cfg.seed, + train_collator=self._train_collator, + ) + + return ( + dict.fromkeys(_ALL_MODALITY_KEYS) | {TEXT_DECODER: calib_loader}, + dict.fromkeys(_ALL_MODALITY_KEYS) | {TEXT_DECODER: train_loader}, + dict.fromkeys(_ALL_MODALITY_KEYS) | {TEXT_DECODER: val_loader}, + ) + def build_runtime_dataloader( self, message: MessageSample, @@ -319,6 +516,109 @@ def build_runtime_dataloader( for modality, dataset in datasets.items() } + @staticmethod + def _split_train_val( + dataset: Dataset, + train_val_ratio: float, + batch_size: int, + seed: int, + train_collator, + ) -> Tuple[DataLoader, Optional[DataLoader]]: + n_total = len(dataset) + n_train = max(1, int(n_total * train_val_ratio)) + n_val = n_total - n_train + + generator = torch.Generator().manual_seed(seed) + indices = torch.randperm(n_total, generator=generator).tolist() + + train_loader = DataLoader( + Subset(dataset, indices[:n_train]), + batch_size=batch_size, + shuffle=True, + drop_last=True, + collate_fn=train_collator, + ) + val_loader: Optional[DataLoader] = ( + DataLoader( + Subset(dataset, indices[n_train:]), + batch_size=batch_size, + shuffle=False, + drop_last=True, + collate_fn=train_collator, + ) + if n_val > 0 + else None + ) + logging.info( + "QAT train/val split (seed=%d): total=%d train=%d val=%d", + seed, + n_total, + n_train, + n_val, + ) + return train_loader, val_loader + + @staticmethod + def _split_qat_dataset( + dataset: Dataset, + calib_train_ratio: float, + train_val_ratio: float, + batch_size: int, + seed: int, + calib_collator, + train_collator, + ) -> Tuple[DataLoader, DataLoader, Optional[DataLoader]]: + n_total = len(dataset) + n_calib = max(1, int(n_total * calib_train_ratio)) + n_remaining = n_total - n_calib + n_train = max(1, int(n_remaining * train_val_ratio)) + n_val = n_remaining - n_train + + generator = torch.Generator().manual_seed(seed) + indices = torch.randperm(n_total, generator=generator).tolist() + + calib_idx = indices[:n_calib] + train_idx = indices[n_calib : n_calib + n_train] + val_idx = indices[n_calib + n_train : n_calib + n_train + n_val] + + calib_loader = DataLoader( + Subset(dataset, calib_idx), + batch_size=batch_size, + shuffle=False, + drop_last=batch_size > 1, + collate_fn=calib_collator, + ) + # drop_last=True: the exported graph module has static input shapes, so every batch + # must have exactly batch_size samples. Incomplete trailing batches are dropped. + train_loader = DataLoader( + Subset(dataset, train_idx), + batch_size=batch_size, + shuffle=True, + drop_last=True, + collate_fn=train_collator, + ) + val_loader: Optional[DataLoader] = ( + DataLoader( + Subset(dataset, val_idx), + batch_size=batch_size, + shuffle=False, + drop_last=True, + collate_fn=train_collator, + ) + if n_val > 0 + else None + ) + logging.info( + "QAT splits (seed=%d): total=%d calib=%d (%.0f%%) train=%d val=%d", + seed, + n_total, + n_calib, + calib_train_ratio * 100, + n_train, + n_val, + ) + return calib_loader, train_loader, val_loader + @staticmethod def _log_dataset_stats( datasets: Dict[str, Dataset], diff --git a/examples/qualcomm/oss_scripts/llama/dataset/collators.py b/examples/qualcomm/oss_scripts/llama/dataset/collators.py index 439718b7c12..319ea6e3aa7 100644 --- a/examples/qualcomm/oss_scripts/llama/dataset/collators.py +++ b/examples/qualcomm/oss_scripts/llama/dataset/collators.py @@ -4,10 +4,15 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging from typing import Dict, List, Tuple, Union import torch import torch.nn.functional as F +from executorch.examples.qualcomm.oss_scripts.llama.dataset.targets import ( + make_causal_labels, + make_conversation_labels, +) from executorch.examples.qualcomm.oss_scripts.llama.masking_utils import AttentionMask @@ -73,3 +78,50 @@ def __call__(self, batch: List[Tuple]) -> Dict[str, Union[torch.Tensor, Tuple]]: "labels": None, "token_ids": token_ids, } + + +class LLMTrainingCollator(LLMCalibCollator): + """Collator for QAT training batches — pads tokens, builds masks, and appends labels. + + When assistant_masks are present (from HF chat datasets), only assistant-turn + positions are supervised. Otherwise causal next-token labels are used. + """ + + _warned_empty_assistant_mask = False + + def __call__( + self, batch: List[Tuple] + ) -> Dict[str, Union[torch.Tensor, Tuple, List]]: + model_inputs = super().__call__(batch) + assistant_masks = [item[1] for item in batch] + + model_inputs["labels"] = torch.stack( + [ + torch.tensor( + self._make_labels(seq, mask), + dtype=torch.int64, + ) + for seq, mask in zip(model_inputs["token_ids"], assistant_masks) + ] + ) + + return model_inputs + + def _make_labels(self, tokens: List[int], assistant_mask: List[int]) -> List[int]: + """Assistant-only labels, warning once when a mask supervises nothing. + + A None or all-zero mask means the sample has no assistant turns, so + make_conversation_labels falls back to full causal next-token labels. + """ + if assistant_mask is None or not any(assistant_mask): + if not LLMTrainingCollator._warned_empty_assistant_mask: + logging.warning( + "assistant_mask is empty or all zeros (plain-text corpus, or " + "chat template missing `{%% generation %%}` markers); falling " + "back to causal next-token labels. Supervision will cover the " + "full sequence, not just assistant turns." + ) + LLMTrainingCollator._warned_empty_assistant_mask = True + return make_causal_labels(tokens, self._max_context_len) + + return make_conversation_labels(tokens, assistant_mask, self._max_context_len) diff --git a/examples/qualcomm/oss_scripts/llama/dataset/config.py b/examples/qualcomm/oss_scripts/llama/dataset/config.py index c8794209145..f412f33cc89 100644 --- a/examples/qualcomm/oss_scripts/llama/dataset/config.py +++ b/examples/qualcomm/oss_scripts/llama/dataset/config.py @@ -17,12 +17,39 @@ class DataConfig: max_context_len: int token_dtype: Optional[str] = None - # calibration + batch_size: int = 1 + seed: int = 42 + + # PTQ calibration (also used as QAT Mode 2 calib source) calib_tasks: Optional[List[str]] = None calib_limit: int = 1 calib_num_fewshot: Optional[int] = None calib_samples: Optional[List[str]] = field(default_factory=list) - batch_size: int = 1 + calib_hf_dataset: Optional[str] = None + calib_hf_limit: int = 1 + + # QAT Mode 1 — full split (single dataset pool, auto-split into calib/train/val) + qat_full_tasks: Optional[List[str]] = None + qat_full_hf_dataset: Optional[str] = None + qat_full_hf_limit: int = 200 + qat_full_limit: int = 200 + calib_train_ratio: float = 0.2 # fraction of full pool for calib; rest → train+val + + # QAT Mode 2 — explicit split (separate calib and train sources) + train_tasks: Optional[List[str]] = None + train_limit: int = 180 + train_hf_dataset: Optional[str] = None + train_hf_limit: int = 180 + + # Shared by Mode 1 & 2: fraction of the train+val pool used for training + train_val_ratio: float = 1.0 # 1.0 disables validation + + @property + def qat_mode(self) -> str: + """'full' if a single pool is provided for auto-splitting; 'explicit' otherwise.""" + if self.qat_full_tasks or self.qat_full_hf_dataset: + return "full" + return "explicit" @classmethod def from_args(cls, args: argparse.Namespace) -> "DataConfig": @@ -33,9 +60,21 @@ def from_args(cls, args: argparse.Namespace) -> "DataConfig": if getattr(args, "embedding_quantize", None) else torch.int32 ), + batch_size=getattr(args, "batch_size", 1), calib_tasks=getattr(args, "calib_tasks", None), calib_limit=getattr(args, "calib_limit", 1), calib_num_fewshot=getattr(args, "calib_num_fewshot", None), calib_samples=getattr(args, "calib_samples", None) or [], - batch_size=getattr(args, "batch_size", 1), + calib_hf_dataset=getattr(args, "calib_hf_dataset", None), + calib_hf_limit=getattr(args, "calib_hf_limit", 1), + qat_full_tasks=getattr(args, "qat_full_tasks", None), + qat_full_hf_dataset=getattr(args, "qat_full_hf_dataset", None), + qat_full_hf_limit=getattr(args, "qat_full_hf_limit", 200), + qat_full_limit=getattr(args, "qat_full_limit", 200), + calib_train_ratio=args.calib_train_ratio, + train_tasks=getattr(args, "train_tasks", None), + train_limit=getattr(args, "train_limit", 180), + train_hf_dataset=getattr(args, "train_hf_dataset", None), + train_hf_limit=getattr(args, "train_hf_limit", 180), + train_val_ratio=args.train_val_ratio, ) diff --git a/examples/qualcomm/oss_scripts/llama/dataset/constants.py b/examples/qualcomm/oss_scripts/llama/dataset/constants.py new file mode 100644 index 00000000000..fa9c72ba1d3 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/dataset/constants.py @@ -0,0 +1,7 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +LABEL_IGNORE_INDEX: int = -100 diff --git a/examples/qualcomm/oss_scripts/llama/dataset/datasets.py b/examples/qualcomm/oss_scripts/llama/dataset/datasets.py index 66f52a35d7d..b939097ad25 100644 --- a/examples/qualcomm/oss_scripts/llama/dataset/datasets.py +++ b/examples/qualcomm/oss_scripts/llama/dataset/datasets.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from typing import List, Tuple +from typing import List, Optional, Tuple from torch.utils.data import Dataset @@ -12,20 +12,27 @@ class LLMDataset(Dataset): """TEXT_DECODER dataset — pure token storage, no padding or masking. - sequences: raw (unpadded) token-id lists, one per sample. + sequences: raw (unpadded) token-id lists, one per sample. + assistant_masks: parallel token-level 0/1 mask (1 = assistant turn); + None for plain-text tasks (causal labels used instead). """ def __init__( self, sequences: List[List[int]], + assistant_masks: Optional[List[List[int]]] = None, ): self.sequences = sequences + self.assistant_masks = assistant_masks def __len__(self) -> int: return len(self.sequences) - def __getitem__(self, idx: int) -> Tuple[List[int]]: - return (self.sequences[idx],) + def __getitem__(self, idx: int) -> Tuple[List[int], Optional[List[int]]]: + assistant_mask = ( + self.assistant_masks[idx] if self.assistant_masks is not None else None + ) + return self.sequences[idx], assistant_mask class ModalityEncoderDataset(Dataset): diff --git a/examples/qualcomm/oss_scripts/llama/dataset/loaders.py b/examples/qualcomm/oss_scripts/llama/dataset/loaders.py index 97c8868168d..468bc52132f 100644 --- a/examples/qualcomm/oss_scripts/llama/dataset/loaders.py +++ b/examples/qualcomm/oss_scripts/llama/dataset/loaders.py @@ -9,10 +9,11 @@ import logging import os from pathlib import Path -from typing import List, Optional, Union +from typing import List, Optional, Tuple, Union import requests import torch +from datasets import load_dataset as hf_load_dataset from executorch.examples.qualcomm.oss_scripts.llama.dataset.schema import MessageSample from huggingface_hub import hf_hub_download from pytorch_tokenizers.hf_tokenizer import HuggingFaceTokenizer @@ -103,3 +104,89 @@ def _model_call(self, inps): tasks_limit, ) return collector.sequences + + +def _assistant_mask_from_boundaries( + messages: List[dict], + full_tokens: List[int], + tokenizer_wrapper, +) -> List[int]: + """Token-level assistant mask computed by re-tokenizing message prefixes. + + Fallback for chat templates that lack the `{% generation %}` / + `{% endgeneration %}` markers HuggingFace's `return_assistant_tokens_mask` + relies on. Instead of trusting the template annotation, we recover each turn's + token span from the length of tokenize(messages[:i+1]) and mark assistant spans. + The prefix tokenization must match the full tokenization (same template, + add_generation_prompt=False) for boundaries to line up. + """ + mask = [0] * len(full_tokens) + boundaries = [0] + for i in range(len(messages)): + prefix = tokenizer_wrapper.chat_template( + messages[: i + 1], + tokenize=True, + add_generation_prompt=False, + return_dict=False, + ) + boundaries.append(len(prefix)) + + for msg, start, end in zip(messages, boundaries, boundaries[1:]): + if msg.get("role") == "assistant": + for j in range(start, min(end, len(full_tokens))): + mask[j] = 1 + return mask + + +def load_hf_chat_dataset( + dataset_name: str, + tokenizer_wrapper, + max_context_len: int, + num_samples: int = 1000, +) -> Tuple[List[List[int]], List[List[int]]]: + """Tokenize a HuggingFace chat dataset (requires a 'messages' field). + + Returns (token_sequences, assistant_masks_list). + assistant_masks_list[i] is a token-level 0/1 mask: 1 = assistant turn. + Samples that fail tokenization or exceed max_context_len are dropped. + """ + + ds = hf_load_dataset(dataset_name, split="train") + ds = ds.select(range(min(num_samples, len(ds)))) + + input_sequences: List[List[int]] = [] + assistant_masks_list: List[List[int]] = [] + + for sample in ds: + messages = sample.get("messages", []) + full_tokens = [] + if not messages: + continue + try: + result = tokenizer_wrapper.chat_template( + messages, + tokenize=True, + add_generation_prompt=False, + return_assistant_tokens_mask=True, + return_dict=True, + ) + full_tokens = result["input_ids"] + assistant_masks = result.get("assistant_masks", [0] * len(full_tokens)) + if not any(assistant_masks): + assistant_masks = _assistant_mask_from_boundaries( + messages, full_tokens, tokenizer_wrapper + ) + except Exception: + continue + if not full_tokens or len(full_tokens) > max_context_len: + continue + input_sequences.append(full_tokens) + assistant_masks_list.append(assistant_masks) + + logging.info( + "%s: loaded %d sequences (requested %d)", + dataset_name, + len(input_sequences), + num_samples, + ) + return input_sequences, assistant_masks_list diff --git a/examples/qualcomm/oss_scripts/llama/dataset/targets.py b/examples/qualcomm/oss_scripts/llama/dataset/targets.py new file mode 100644 index 00000000000..8781dd59936 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/dataset/targets.py @@ -0,0 +1,43 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import List + +from executorch.examples.qualcomm.oss_scripts.llama.dataset.constants import ( + LABEL_IGNORE_INDEX, +) + + +def make_causal_labels(tokens: List[int], max_context_len: int) -> List[int]: + """ + Shifted next-token labels: labels[i] = tokens[i+1]; padding positions will be ignored. + + """ + n = min(len(tokens) - 1, max_context_len - 1) + labels = [LABEL_IGNORE_INDEX] * max_context_len + if n > 0: + labels[:n] = tokens[1 : n + 1] + return labels + + +def make_conversation_labels( + tokens: List[int], + assistant_mask: List[int], + max_context_len: int, +) -> List[int]: + """ + Conversation labels: only assistant-turn positions supervised; others are ignored. + + Same pre-shift convention as make_causal_labels. assistant_mask[i] == 1 marks + token i as an assistant-response token. Callers must guarantee the mask supervises + at least one token; an all-zero mask yields all-ignored labels here. + """ + n = min(len(tokens), max_context_len) + labels = [LABEL_IGNORE_INDEX] * max_context_len + for i in range(n - 1): + if i < len(assistant_mask) and assistant_mask[i]: + labels[i - 1] = tokens[i] + return labels diff --git a/examples/qualcomm/oss_scripts/llama/inference/decoder.py b/examples/qualcomm/oss_scripts/llama/inference/decoder.py index 18006af3f61..03b38e90a00 100644 --- a/examples/qualcomm/oss_scripts/llama/inference/decoder.py +++ b/examples/qualcomm/oss_scripts/llama/inference/decoder.py @@ -146,3 +146,19 @@ def predict_step( *self.get_inputs(input_ids, attn_mask, hidden_states, tok_embedding) ) return logits + + def make_forward_fn(self) -> Callable: + k = self._k_caches + v = self._v_caches + pos = self.pos_ids + + def _forward( + model: Union[torch.nn.Module, torch.fx.GraphModule], + model_inputs: dict, + ) -> torch.Tensor: + input_ids = model_inputs["input_ids"] + attn_masks = model_inputs["attention_mask"] + out = model(input_ids, *attn_masks, pos, *k, *v) + return out[0] if isinstance(out, (tuple, list)) else out + + return _forward diff --git a/examples/qualcomm/oss_scripts/llama/llama.py b/examples/qualcomm/oss_scripts/llama/llama.py index edf719a6fe5..73a3385016b 100755 --- a/examples/qualcomm/oss_scripts/llama/llama.py +++ b/examples/qualcomm/oss_scripts/llama/llama.py @@ -525,6 +525,7 @@ def _build_parser(): default=8, type=int, ) + parser.add_argument( "--use_attention_sink", default=None, @@ -620,6 +621,21 @@ def _build_parser(): "Multiple files are merged.", ) + parser.add_argument( + "--calib_hf_dataset", + type=str, + default=None, + help="HuggingFace chat dataset to use as additional calibration data " + "(e.g. 'HuggingFaceTB/smol-smoltalk').", + ) + + parser.add_argument( + "--calib_hf_limit", + type=int, + default=1, + help="Number of samples to load from --calib_hf_dataset.", + ) + parser.add_argument( "--batch_size", type=int, @@ -629,6 +645,119 @@ def _build_parser(): "PREFILL graphs always use batch size 1.", ) + parser.add_argument( + "--qat", + action="store_true", + help="Enable Quantization-Aware Training (QAT). If not set, defaults to PTQ.", + ) + + parser.add_argument( + "--train_config", + type=str, + default=os.path.join(os.path.dirname(__file__), "train", "config", "qad.yaml"), + help="(QAT) Path to a YAML file overriding TrainingArgs defaults " + "(train/config/config.py) — e.g. epochs, lr, alpha, temperature, " + "grad_accum_steps, max_grad_norm, warmup_ratio. Defaults to " + "train/config/qad.yaml.", + ) + + parser.add_argument( + "--lr_config", + type=str, + default=os.path.join( + os.path.dirname(__file__), "train", "config", "lr_config.yaml" + ), + help="(QAT) Path to a YAML file mapping glob patterns over " + "model.named_parameters() names to per-param-group optimizer kwargs " + "(lr, weight_decay, betas, eps, ...). Defaults to " + "train/config/lr_config.yaml.", + ) + + parser.add_argument( + "--qat_full_tasks", + nargs="+", + type=str, + default=None, + help="(QAT) list of lm-eluther tasks for calib and training data." + "Usage: --qat_full_tasks task1 task2" + "The full tasks is split into calib/train/val by --calib_train_ratio / --train_val_ratio. ", + ) + + parser.add_argument( + "--qat_full_limit", + type=int, + default=200, + help="(QAT) number of samples for full task limits.", + ) + + parser.add_argument( + "--qat_full_hf_dataset", + type=str, + default=None, + help="(QAT) HuggingFace instruct dataset name for calib and training " + "(e.g. 'HuggingFaceTB/smol-smoltalk')." + "The full HuggingFace dataset is split into calib/train/val by --calib_train_ratio / --train_val_ratio. ", + ) + + parser.add_argument( + "--qat_full_hf_limit", + type=int, + default=200, + help="(QAT) Number of samples to load from --qat_full_hf_dataset.", + ) + + parser.add_argument( + "--calib_train_ratio", + type=float, + default=0.2, + help="(QAT) Fraction of the full data pool for PTQ calibration. " + "The remaining (1 - calib_train_ratio) goes to train + val. " + "e.g. 0.2 → 20%% calib, 80%% for train+val.", + ) + + parser.add_argument( + "--train_tasks", + nargs="+", + type=str, + default=None, + help="(QAT) list of lm-eluther tasks for training data. " + "Usage: --train_tasks task1 task2." + "And you need to use --calib_tasks to specify calibration data separately. ", + ) + + parser.add_argument( + "--train_limit", + type=int, + default=1, + help="(QAT) number of samples for train task limits.", + ) + + parser.add_argument( + "--train_hf_dataset", + type=str, + default=None, + help="(QAT) HuggingFace instruct dataset for training " + "(e.g. 'HuggingFaceTB/smol-smoltalk')." + "And you may need to use --calib_tasks to specify calibration data separately. ", + ) + + parser.add_argument( + "--train_hf_limit", + type=int, + default=1000, + help="(QAT) Number of samples to load from --train_hf_dataset. " + "And you may need to use --calib_hf_dataset to specify calibration data separately. ", + ) + + parser.add_argument( + "--train_val_ratio", + type=float, + default=1.0, + help="(QAT) Fraction of the non-calib samples used for training. " + "The remaining becomes the validation set. " + "e.g. 0.9 → 90%% train, 10%% val. 1.0 disables validation.", + ) + parser.add_argument( "-F", "--use_fp16", @@ -639,6 +768,16 @@ def _build_parser(): parser.add_argument("-v", "--verbose", action="store_true") + parser.add_argument( + "--freeze_all_params", + action="store_true", + help="(QAT CI only) Freeze all model weight parameters during QAT training so that " + "only quantization scale/zero_point are updated via observer calibration — weights are not " + "tuned. This is used as a controlled baseline to isolate the effect of weight updates: " + "comparing --freeze_all_params vs normal QAT shows how much benefit comes from weight " + "adaptation alone.", + ) + parser.add_argument( "--quant_recipe_suggestion", action="store_true", @@ -744,18 +883,45 @@ def export_llama(args) -> None: ) if args.eval_tasks is not None: raise ValueError("Multimodal models do not support --eval_tasks.") + if args.qat: + raise ValueError("QAT is not supported for multimodal models (VLM/ALM).") if not args.pre_gen_pte: - if is_multimodal and args.calib_samples is None: - raise ValueError( - "For MLLMs calibration data is required for compilation. " - "Provide --calib_samples with a vision/audio JSON file." - ) - if not is_multimodal and not any((args.calib_tasks, args.calib_samples)): - raise ValueError( - "For LLMs calibration data is required for compilation. " - "Provide --calib_tasks or --calib_samples." + if is_multimodal: + if args.calib_samples is None: + raise ValueError( + "For MLLMs calibration data is required for compilation. " + "Provide --calib_samples with a vision/audio JSON file." + ) + else: + has_calib = any( + (args.calib_tasks, args.calib_samples, args.calib_hf_dataset) ) + has_full = args.qat_full_tasks or args.qat_full_hf_dataset + + if args.qat: + has_train = args.train_tasks or args.train_hf_dataset + if has_full and has_train: + raise ValueError( + "QAT accepts only one training-data source: use " + "--qat_full_tasks / --qat_full_hf_dataset for auto-split into " + "calib/train/val, or --train_tasks / --train_hf_dataset paired " + "with --calib_* for explicit split — not both." + ) + if not has_full and not has_train: + raise ValueError( + "QAT requires training data: provide --qat_full_tasks / " + "--qat_full_hf_dataset (auto-split), or --train_tasks / " + "--train_hf_dataset together with --calib_* (explicit split)." + ) + + # Calibration data is required, except in QAT auto-split mode where it + # is carved from the --qat_full_* pool. + if not (args.qat and has_full) and not has_calib: + raise ValueError( + "For LLMs calibration data is required for compilation. " + "Provide --calib_tasks or --calib_samples or --calib_hf_dataset." + ) if args.pre_gen_pte: text_decoder_pte_path = f"{args.pre_gen_pte}/{pte_filenames[TEXT_DECODER]}.pte" diff --git a/examples/qualcomm/oss_scripts/llama/quantization_guidance.md b/examples/qualcomm/oss_scripts/llama/quantization_guidance.md index f42f5be38f0..7d882c7595c 100644 --- a/examples/qualcomm/oss_scripts/llama/quantization_guidance.md +++ b/examples/qualcomm/oss_scripts/llama/quantization_guidance.md @@ -1,4 +1,221 @@ -# LLMs Quantization Guidance +# LLM Quantization Guide + +End-to-end quantization for Large Language model (LLMs), Vision-Language models (VLMs), and +Audio-Language models (ALMs) on the Qualcomm HTP backend. This guide covers both how to +run quantization and how to quantize an LLM well (the recommended workflow). + +> **Recommendation:** PTQ with mixed precision (W8 per-channel + W4 per-block quantization) +> usually works well enough on its own. If your target is more aggressive: W4 or W2 +> per-channel, please consider QAT instead. + +Everything below runs through `llama.py`. + +## Support + +- **LLM** — PTQ and QAT. +- **VLM** — PTQ only (QAT not yet supported). +- **ALM** — PTQ only (QAT not yet supported). + +--- + +# Usage + +## End-to-end PTQ (Post-Training Quantization) +```bash +python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} --decoder_model smollm2_135m --model_mode hybrid --prefill_ar_len 128 --max_seq_len 1024 --prompt "I would like to learn python, could you teach me with a simple example?" --calib_tasks wikitext --calib_limit 1 +``` + +Calibration data only affects quantization quality — it does not influence the inference +prompt or the runtime evaluation dataset. See [Data flags reference](#data-flags-reference). + +## End-to-end QAT (Quantization-Aware Training) + +QAT (Quantization-Aware Training) trains the model with fake-quantization nodes active so +weights learn to compensate for quantization error before the graph is lowered to QNN. +Add `--qat` and a training-data source. By default QAT uses **Quantization-Aware +Distillation** (QAD): a frozen FP32 copy of the model acts as a teacher and the +fake-quantized student learns to match its logits. + +```bash +python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} --decoder_model smollm2_135m --model_mode hybrid --max_seq_len 1024 --prompt "What does 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 --train_val_ratio 1.0 --train_config ./examples/qualcomm/oss_scripts/llama/train/config/qad.yaml --lr_config ./examples/qualcomm/oss_scripts/llama/train/config/lr_config.yaml +``` + +QAT needs both a **calibration set** (to seed observer statistics, via `--calib_*`) and a +**training set** (to adapt weights, via `--train_*`), as in the example above. See the +[training-data table](#training-data-qat-only) for all the flags. Recipe authoring and +hyperparameters are covered in +[Quantization-Aware Training / Distillation](#quantization-aware-training--distillation). + +## Data flags reference + +Calibration, training, and runtime evaluation use separate flag sets and can target +different tasks or limits as needed. + +### Calibration data + +Required for compilation. Provide at least one source; multiple sources are concatenated. + +| Flag | Description | +|---|---| +| `--calib_tasks` | One or more lm_eval tasks (LLM-only). | +| `--calib_limit` | Number of samples per task (default 1). | +| `--calib_num_fewshot` | Few-shot examples per calibration sample. | +| `--calib_samples` | Custom conversation JSON files (see format below). Required for VLM/ALM. | +| `--calib_hf_dataset` | A HuggingFace chat dataset name as additional calibration data. | +| `--calib_hf_limit` | Number of samples to load from `--calib_hf_dataset` (default 1). | + +For LLMs, provide at least one calibration source. **VLM and ALM models can only be +calibrated through `--calib_samples`** (a JSON file with the media inputs); the +lm_eval-based (`--calib_tasks`) and HuggingFace-dataset (`--calib_hf_dataset`) sources are +LLM-only. + +### Training data + +QAT reuses the `--calib_*` flags above for the calibration set, and takes a separate +training set from the flags below. An optional validation split is carved from the +training pool via `--train_val_ratio`. + +| Flag | Description | +|---|---| +| `--train_tasks` | lm_eval tasks for training data. | +| `--train_limit` | Number of samples for train tasks (default 1). | +| `--train_hf_dataset` | HuggingFace instruct dataset for training. | +| `--train_hf_limit` | Number of samples from `--train_hf_dataset` (default 1000). | +| `--train_val_ratio` | Fraction of the training pool used for training; rest → val. `1.0` disables val (default 1.0). | + +When a HuggingFace instruct dataset is used, labels apply assistant-only loss — only +assistant-turn tokens contribute to the loss; system and user turns are masked. Plain-text +tasks (e.g. wikitext) fall back to standard causal next-token labels. + +### Evaluation data + +Runtime/offline evaluation, independent from calibration. + +| Flag | Description | +|---|---| +| `--eval_methods` | One or more of `prompt_eval`, `tasks_eval`, `sqnr_eval` (default `prompt_eval`). | +| `--eval_tasks` | lm_eval tasks to evaluate (for `tasks_eval`). | +| `--eval_limit` | Number of samples to evaluate (default 1). | +| `--eval_num_fewshot` | Few-shot examples per evaluation sample. | + +### Custom calibration samples (`--calib_samples`) + +`--calib_samples` accepts one or more JSON files. Each file is a flat list of sample +objects. Each sample has a `messages` field following the HuggingFace chat template, and +an optional `files` field for media inputs (local paths or URLs): + +```json +[ + { + "files": ["path/or/url/to/files"], + "messages": [ + {"role": "user", "content": "..." }, + {"role": "assistant", "content": "..."} + ] + } +] +``` + +`files` is only required for multimodal models (VLM: image paths/URLs, ALM: audio +paths/URLs). For LLM-only models, `files` can be omitted. `content` can be a plain string +or a list of HuggingFace content blocks (e.g. `[{"type": "image"}, {"type": "text", +"text": "..."}]` for vision inputs). + +Ready-to-use examples for each model type are provided under `assets/samples/`: + +| Model type | Example file | +|---|---| +| LLM | [assets/samples/text.json](assets/samples/text.json) | +| ALM (audio) | [assets/samples/audio.json](assets/samples/audio.json) | +| VLM (vision) | [assets/samples/vision.json](assets/samples/vision.json) | + +--- + +# Quantization Workflow + +Treat quantization as an iterative optimization loop, not a one-shot conversion. Start +with a simple baseline, refine with recipe-driven mixed precision, use SQNR to locate the +quantization bottlenecks, and escalate to QAT only when recipe refinement can no longer +close the gap. + +
+ End-to-end quantization decision workflow +
+ +| Technique | Primary goal | Description | +|---|---|---| +| **PTQ** | Produce a quantized model quickly. | With a uniform `16a8w` precision, accuracy usually stays on par with the FP32/FP16 model. Use it as the baseline. | +| **Mixed Precision** | Per-layer / per-component precision control. | Use a Quant Recipe to assign precision to each layer or component individually. The simplest and most effective way to fix accuracy issues, and the most balanced trade-off. | +| **SQNR Analysis** | Measure per-layer quantization sensitivity. | When quantizing at `16a8w` or LPBQ (`16a4w_block`), quickly locates where the accuracy bottlenecks are and auto-generates better recipes. | +| **QAT / Distillation** | Recover accuracy by training. | For aggressive targets like per-channel W4/W2, where PTQ alone cannot reach acceptable accuracy, training lets the weights adapt to quantization noise and recover the quality. | + +- **Protect the LM head carefully** — it directly controls token logits and has an + outsized impact on perplexity. +- **Keep the data distribution aligned** across calibration, QAT, and runtime prompts. + + +## PTQ + +The quickest path to quantize a model. With `16a8w` precision, accuracy usually stays on +par with the FP32/FP16 model, so it makes a solid baseline. + +1. Start from the trained FP32 model. +2. Set a baseline precision: `16a8w` is the recommended starting point. +3. Run calibration with representative calibration data. +4. Export and evaluate perplexity, token generation quality, and on-device performance. +5. If quality is insufficient, move to mixed precision. + +## Quant Recipe Mixed Precision + +A `StaticLLMQuantRecipe` declaratively assigns different quantization precision to +different parts of the graph, so you can promote a few sensitive layers to higher +precision while keeping the rest aggressively compressed. This is most useful for W4/W8 +mixed precision. + +```python +class Llama3_1BQuantRecipe(StaticLLMQuantRecipe): + default_quant_dtype = QuantDtype.use_16a4w_block + + def __init__(self, verbose: bool = False): + super().__init__() + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + ) + .add_node_target( + {torch.ops.aten.conv2d.default}, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 32, 1, 1)}, + ) + .add_regex( + {r"output\.conv", r"layers\.[0-3]\.feed_forward\.w2_conv"}, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + ) + ) +``` + +The default keeps the whole graph at `16a4w_block`. `add_regex` then promotes the LM head +(`output.conv`) and the early down-projection layers (`layers.[0-3].feed_forward.w2_conv`) +to `16a8w`, since those are the layers most sensitive to quantization error. + +Targeting is by regex or graph node target: + +| Pattern | Meaning | +|---|---| +| `default_quant_dtype = use_16a4w_block` | LPBQ 4-bit block-wise weights for everything not overridden. | +| `add_regex({r"output\.conv"}, use_16a8w, ...)` | Promote the LM head to 8-bit weights. | +| `add_regex({r"layers\..*\.attention\.wv.*"}, use_8a4w, ...)` | 8-bit activations on value projections. | + ## Mixed-Precision Quantization with SQNR Analysis @@ -6,6 +223,13 @@ When deploying LLMs at low precision (for example, `16a4w_block`), some layers c mix_precision_analyzer.py is an analysis tool that helps you identify these quantization-sensitive layers and provides a directional starting point for mixed-precision tuning. It lets you selectively upgrade only the most quantization-sensitive layers to higher precision, while keeping the rest of the model at the aggressive baseline (e.g. 16a4w_block). The tool does not aim to find a globally optimal quantization recipe, but it helps narrow the search space so you can iterate from a directional starting point rather than guessing. +
+ SQNR sensitivity analysis design and pipeline +
+ +The pipeline runs in three phases: **(1) Quantize**, **(2) Compare** per-layer SQNR against +the FP32 model, and **(3) Suggest** a refined recipe. Details below. + ### Overview `mix_precision_analyzer.py` provides two classes and one module-level function: @@ -129,3 +353,104 @@ from executorch.examples.qualcomm.oss_scripts.llama.static_llm_quant_recipe impo - Once satisfied, copy the final recipe into `static_llm_quant_recipe.py` as the permanent recipe for the model. > **Note:** The primary purpose of this SQNR analysis is just to provide a guiding direction for mixed-precision quantization. While it identifies which layers are most sensitive and suggests a reasonable mixed-precision combination, it does not guarantee the best possible combination for every model. Please note that not every model will truly benefit from this mixed precision analysis, as the overall effectiveness of PTQ (Post-Training Quantization) can still be limited. The generated recipe classes are a starting point for exploration, not a final answer. You may need to experiment with different block sizes, different threshold values, or manually crafting overrides for specific layer groups to find the optimal accuracy/performance trade-off for your target model. If your requirement is to push the performance to the extreme limits, please try QAT (Quantization-Aware Training) instead. + + +## Quantization-Aware Training / Distillation + +Use QAT when PTQ + mixed precision can no longer keep the model at acceptable accuracy +under aggressive low-bit quantization (e.g. per-channel W2/W4). QAT lets the weights adapt to +quantization error during training, recovering quality that PTQ cannot. + + +
+ QAT phases: prepare, calibrate, train, deploy +
+ +**Phase 1 — Prepare.** `prepare_qat_pt2e` inserts `FakeQuantize` modules and applies the +QAT recipe (which layers are quantized, trainable, or frozen). + +**Phase 2 — Calibrate.** An observer pass runs with fake-quant disabled to seed sensible +scale/zp, identical to the PTQ path. Optional SeqMSE can refine these initial encodings. + +**Phase 3 — Train (QAD).** With encodings seeded, observers are disabled and fake-quant is +enabled, so every forward pass now runs with quantization noise. The loss compares the +fake-quantized student's output against the FP32 teacher's, and gradients from that loss +update the weights to compensate for the quantization noise. The default loss is +Quantization-Aware Distillation: + +``` +loss = alpha * KL(teacher ‖ student) + (1 - alpha) * CE(student, labels) +``` + +The KL term aligns the quantized student to the FP32 teacher's distribution; the CE term +keeps it anchored to the labels. Default `alpha = 0.99` (KD-dominant) — logits carry more +information per token than hard labels, so distillation is weighted more heavily. + +**Phase 4 — Deploy.** Compare the quantized model against the FP32 model — evaluation +metrics (e.g. PPL, MMLU) and token generation output — to confirm QAT actually recovered +the accuracy PTQ lost. + +Training hyperparameters (`epochs`, `lr`, `alpha`, `temperature`, `grad_accum_steps`, +`max_grad_norm`) come from `TrainingArgs` defaults in `train/config/config.py`. Override +any subset with `--train_config path/to/config.yaml` (see `train/config/qad.yaml`). + +### Defining a QAT recipe + +PTQ recipes should not be reused for QAT. Depending on the model, some parameters may need to be frozen during training to keep an +already-sensitive or already-stable component. Create a +`StaticLLMQATRecipe` subclass and register it as `qat_recipe` on the model config: + +```python +# static_llm_quant_recipe.py +class MyModelQATRecipe(StaticLLMQATRecipe): + default_quant_dtype = QuantDtype.use_16a4w + frozen_param_patterns = [] + + def __init__(self, verbose=False): + super().__init__() + self.recipe = QuantRecipe( + self.default_quant_dtype, + is_qat=True, + act_observer=MovingAverageMinMaxObserver, + ) +``` + +```python +# __init__.py — on the model's @register_llm_model config +@register_llm_model("my_model") +@dataclass(init=False, frozen=True) +class MyModel(LLMModelConfig): + ... + quant_recipe = MyModelQuantRecipe # PTQ recipe + qat_recipe = MyModelQATRecipe # used when --qat is set +``` + +`frozen_param_patterns` is a list of regex strings matched against +`model.named_parameters()`; matched parameters have `requires_grad = False` before the +optimizer is built, so they stay untouched by training and are quantized as plain PTQ +(observer-only, no fake-quant gradient) instead of QAT. For example: + +```python +frozen_param_patterns: List[str] = [ + r"tok_embedding", # Freeze token embeddings to prevent drift in the token space. + r"output\.conv", # Freeze lm head to prevent drift in the token space. +] +``` + +When `--qat` is set and the model has no `qat_recipe`, the pipeline +falls back to its PTQ `quant_recipe`. Currently only `smollm2_135m` ships a registered +`qat_recipe` (`Smollm2QATQuantRecipe`). + + +### QAT tips + +- **Match training data to the inference distribution.** Instruct-formatted / domain data + outperforms a generic corpus when the target is instruction-following. If the training + scenario diverges too far from the runtime scenario (a domain shift), QAT has to learn + domain adaptation on top of quantization recovery, with a fixed dataset size and computation + budget, that split objective leaves less capacity to fix the quantization error itself, + so accuracy recovery suffers. +- **Good PTQ is a prerequisite.** A well-tuned PTQ recipe seeds better encodings and leads + to faster convergence and a lower final loss. +- **Use a small learning rate** Since quantization parameters are often sensitive to even minor updates, a small learning rate is typically recommended for stable convergence. +- **One epoch is usually enough.** QAT is primarily an accuracy-recovery step rather than a full retraining stage. Since the pretrained model is already near a good optimum, the goal is only to compensate for quantization-induced error. As a result, the required parameter updates are typically small, and a single epoch is often sufficient to recover most of the lost accuracy. diff --git a/examples/qualcomm/oss_scripts/llama/quantize/__init__.py b/examples/qualcomm/oss_scripts/llama/quantize/__init__.py index 15a5195282c..b86e388ead2 100644 --- a/examples/qualcomm/oss_scripts/llama/quantize/__init__.py +++ b/examples/qualcomm/oss_scripts/llama/quantize/__init__.py @@ -5,8 +5,9 @@ # LICENSE file in the root directory of this source tree. from executorch.examples.qualcomm.oss_scripts.llama.quantize.ptq import PTQStrategy +from executorch.examples.qualcomm.oss_scripts.llama.quantize.qat import QATStrategy from executorch.examples.qualcomm.oss_scripts.llama.quantize.strategy import ( QuantizationStrategy, ) -__all__ = ["QuantizationStrategy", "PTQStrategy"] +__all__ = ["QuantizationStrategy", "PTQStrategy", "QATStrategy"] diff --git a/examples/qualcomm/oss_scripts/llama/quantize/qat.py b/examples/qualcomm/oss_scripts/llama/quantize/qat.py new file mode 100644 index 00000000000..5747a28d556 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/quantize/qat.py @@ -0,0 +1,105 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Dict, List, Optional + +import torch +from executorch.examples.qualcomm.oss_scripts.llama.inference import DecoderInference +from executorch.examples.qualcomm.oss_scripts.llama.quantize.ptq import PTQStrategy +from executorch.examples.qualcomm.oss_scripts.llama.train.config import TrainingArgs +from executorch.examples.qualcomm.oss_scripts.llama.train.trainer import ( + KDTrainer, + Trainer, +) +from torch.utils.data import DataLoader + + +def _disable_fake_quant_enable_observer(model: torch.fx.GraphModule) -> None: + for _, submodule in model.named_modules(): + if hasattr(submodule, "disable_fake_quant"): + submodule.disable_fake_quant() + if hasattr(submodule, "enable_observer"): + submodule.enable_observer() + + +def _enable_fake_quant_freeze_observer(model: torch.fx.GraphModule) -> None: + for _, submodule in model.named_modules(): + if hasattr(submodule, "enable_fake_quant"): + submodule.enable_fake_quant() + if hasattr(submodule, "disable_observer"): + submodule.disable_observer() + + +class QATStrategy(PTQStrategy): + """Quantization-aware training strategy""" + + def __init__( + self, + inference: DecoderInference, + module: torch.fx.GraphModule, + seq_mse_candidates: int = 0, + tok_embedding: Optional[torch.fx.GraphModule] = None, + ): + super().__init__( + inference=inference, + module=module, + seq_mse_candidates=seq_mse_candidates, + tok_embedding=tok_embedding, + ) + + def _calibrate(self, calib_loader: Dict[str, DataLoader]) -> None: + _disable_fake_quant_enable_observer(self._module) + super()._calibrate(calib_loader) + _enable_fake_quant_freeze_observer(self._module) + + def _train( + self, + training_args: TrainingArgs, + teacher: torch.nn.Module, + train_loader: DataLoader, + val_loader: Optional[DataLoader] = None, + frozen_param_patterns: Optional[List[str]] = None, + ) -> None: + forward_fn = self._inference.make_forward_fn() + if teacher is not None: + KDTrainer( + model=self._module, + teacher=teacher, + args=training_args, + train_loader=train_loader, + forward_fn=forward_fn, + val_loader=val_loader, + frozen_param_patterns=frozen_param_patterns, + ).train() + else: + Trainer( + model=self._module, + args=training_args, + train_loader=train_loader, + forward_fn=forward_fn, + val_loader=val_loader, + frozen_param_patterns=frozen_param_patterns, + ).train() + + def quantize( + self, + calib_loader: Dict[str, DataLoader], + training_args: TrainingArgs = None, + teacher: torch.nn.Module = None, + train_loader: DataLoader = None, + val_loader: Optional[DataLoader] = None, + frozen_param_patterns: Optional[List[str]] = None, + **kwargs, + ) -> torch.fx.GraphModule: + self._calibrate(calib_loader) + self._train( + training_args=training_args, + teacher=teacher, + train_loader=train_loader, + val_loader=val_loader, + frozen_param_patterns=frozen_param_patterns, + ) + return self._module diff --git a/examples/qualcomm/oss_scripts/llama/quantize/strategy.py b/examples/qualcomm/oss_scripts/llama/quantize/strategy.py index 9b74d30efa5..56e9cb5c76f 100644 --- a/examples/qualcomm/oss_scripts/llama/quantize/strategy.py +++ b/examples/qualcomm/oss_scripts/llama/quantize/strategy.py @@ -8,16 +8,16 @@ from typing import Dict, Optional import torch -from executorch.examples.qualcomm.oss_scripts.llama.inference import ModelInference +from executorch.examples.qualcomm.oss_scripts.llama.inference import DecoderInference from torch.utils.data import DataLoader class QuantizationStrategy(ABC): - """Base class for quantization strategies.""" + """Base class for PTQ and QAT quantization strategies.""" def __init__( self, - inference: ModelInference, + inference: DecoderInference, module: torch.fx.GraphModule, tok_embedding: Optional[torch.fx.GraphModule] = None, ): diff --git a/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py b/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py index 4eecf308f88..9598632dab2 100644 --- a/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py +++ b/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from typing import Optional +from typing import List, Optional import torch from executorch.backends.qualcomm.quantizer.custom_annotation import annotate_kv_8bit @@ -13,7 +13,7 @@ QuantRecipe, ) from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype -from torchao.quantization.pt2e import MinMaxObserver +from torchao.quantization.pt2e import MinMaxObserver, MovingAverageMinMaxObserver class StaticLLMQuantRecipe: @@ -46,6 +46,12 @@ def get_logits_output_bit_width(self) -> int: return 32 if self.default_quant_dtype is None else 16 +class StaticLLMQATRecipe(StaticLLMQuantRecipe): + """Base class for QAT recipes. Adds frozen param patterns.""" + + frozen_param_patterns: List[str] = [] + + class LlamaStories260KQuantRecipe(StaticLLMQuantRecipe): default_quant_dtype = QuantDtype.use_16a4w @@ -666,6 +672,51 @@ def __init__(self, verbose: bool = False): ) +class Smollm2QATQuantRecipe(StaticLLMQATRecipe): + default_quant_dtype = QuantDtype.use_16a8w + frozen_param_patterns: List[str] = [ + r"tok_embedding", # Freeze token embeddings to prevent drift in the token space. + r"output\.conv", # Freeze lm head to prevent drift in the token space. + ] + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + True, + act_observer=MovingAverageMinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + ) + .add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w, + True, + act_observer=MovingAverageMinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + ) + .add_regex( + {r"tok_embeddings"}, + QuantDtype.use_16a8w, + True, + act_observer=MovingAverageMinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + ) + .add_regex( + {r"output\.conv"}, + QuantDtype.use_16a8w, + True, + act_observer=MovingAverageMinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + ) + ) + self.recipe.custom_quant_annotations.append(annotate_kv_8bit) + + class Smollm3QuantRecipe(StaticLLMQuantRecipe): default_quant_dtype = QuantDtype.use_16a4w diff --git a/examples/qualcomm/oss_scripts/llama/train/__init__.py b/examples/qualcomm/oss_scripts/llama/train/__init__.py new file mode 100644 index 00000000000..71c6a95124f --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/train/__init__.py @@ -0,0 +1,25 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.examples.qualcomm.oss_scripts.llama.train.config import TrainingArgs +from executorch.examples.qualcomm.oss_scripts.llama.train.loss import ( + CrossEntropyLoss, + KLDivergenceLoss, +) +from executorch.examples.qualcomm.oss_scripts.llama.train.trainer import ( + BaseTrainer, + KDTrainer, + Trainer, +) + +__all__ = [ + "BaseTrainer", + "CrossEntropyLoss", + "KDTrainer", + "KLDivergenceLoss", + "Trainer", + "TrainingArgs", +] diff --git a/examples/qualcomm/oss_scripts/llama/train/config/__init__.py b/examples/qualcomm/oss_scripts/llama/train/config/__init__.py new file mode 100644 index 00000000000..a421dbf6c05 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/train/config/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.examples.qualcomm.oss_scripts.llama.train.config.config import ( + TrainingArgs, +) + +__all__ = ["TrainingArgs"] diff --git a/examples/qualcomm/oss_scripts/llama/train/config/config.py b/examples/qualcomm/oss_scripts/llama/train/config/config.py new file mode 100644 index 00000000000..1bd84fdfa23 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/train/config/config.py @@ -0,0 +1,45 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import get_type_hints, Optional + +import yaml + + +@dataclass +class TrainingArgs: + """Training algorithm hyperparameters. Model structure and data pipeline concerns live elsewhere.""" + + epochs: int = 1 + lr: float = 2e-5 + alpha: float = 0.99 # KD weight; (1 - alpha) is CE weight + temperature: float = 1.0 + grad_accum_steps: int = 1 + max_grad_norm: float = 1.0 + warmup_ratio: float = 0.05 + # Path to a YAML for per-param-group + # LR (see train/config/lr_config.yaml). Set from --lr_config, not from + # --train_config, so it is not subject to the type conversion below. + lr_config: Optional[str] = None + + @classmethod + def from_yaml(cls, path: str) -> "TrainingArgs": + with open(path) as f: + raw = yaml.safe_load(f) or {} + field_types = get_type_hints(cls) + unknown = raw.keys() - {f.name for f in fields(cls)} + if unknown: + raise ValueError( + f"Unknown TrainingArgs field(s) in {path}: {sorted(unknown)}" + ) + converted = { + k: field_types[k](v) if field_types[k] in (int, float) else v + for k, v in raw.items() + } + return cls(**converted) diff --git a/examples/qualcomm/oss_scripts/llama/train/config/lr_config.yaml b/examples/qualcomm/oss_scripts/llama/train/config/lr_config.yaml new file mode 100644 index 00000000000..b16fca74f6a --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/train/config/lr_config.yaml @@ -0,0 +1,49 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +# +# Pattern names are aligned to the parameter names in +# model/static_llama.py so they match what the trainer sees via +# `model.named_parameters()`: +# tok_embeddings.weight +# layers..attention.{wq,wk,wv,wo}.weight +# layers..feed_forward.{w1,w2,w3}.weight +# layers..attention_norm.weight / layers..ffn_norm.weight +# norm.weight +# output.weight (LM head) +# +# Any keyword accepted by the AdamW optimizer constructor can be specified here. +# Common kwargs for AdamW: +# lr - learning rate +# weight_decay - L2 penalty (overrides the global --weight_decay) +# betas - Adam momentum coefficients [beta1, beta2] +# eps - term added to denominator for numerical stability +# +# Usage: +# --lr_config train/config/lr_config.yaml +# +# Parameters not matched by any pattern below fall back to the global optimizer +# settings (--lr and weight_decay=0.01), so no explicit `default` entry is needed. + +# Output head: lower LR, no weight decay +"*output*": + lr: 4e-6 + weight_decay: 0.0 + +# Attention projections (wq/wk/wv/wo): higher LR + more aggressive momentum. +"*attention.*": + lr: 2e-5 + betas: [0.9, 0.95] + +# Feed-forward projections (w1/w2/w3): higher LR + higher weight decay. +"*feed_forward.*": + lr: 2e-5 + weight_decay: 0.05 + +# Embedding layers: small LR, no weight decay (unfrozen to fine-tune) +"*tok_embeddings*": + lr: 5e-7 + weight_decay: 0.0 + eps: 1e-7 diff --git a/examples/qualcomm/oss_scripts/llama/train/config/qad.yaml b/examples/qualcomm/oss_scripts/llama/train/config/qad.yaml new file mode 100644 index 00000000000..0e4f6d934de --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/train/config/qad.yaml @@ -0,0 +1,20 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Default TrainingArgs (train/config.py) for QAT. Copy this file, edit any +# subset of fields, and pass it via `--train_config path/to/your_config.yaml`. + + +# Distillation +alpha: 0.99 # KD weight; (1 - alpha) is CE weight +temperature: 1.0 + +# Hyperparameters +epochs: 1 # Train epochs +lr: 2.0e-5 +grad_accum_steps: 1 +max_grad_norm: 1.0 +warmup_ratio: 0.05 diff --git a/examples/qualcomm/oss_scripts/llama/train/loss.py b/examples/qualcomm/oss_scripts/llama/train/loss.py new file mode 100644 index 00000000000..4d452a9e6d8 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/train/loss.py @@ -0,0 +1,64 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn.functional as F +from executorch.examples.qualcomm.oss_scripts.llama.dataset.constants import ( + LABEL_IGNORE_INDEX, +) + + +class CrossEntropyLoss: + """Standard next-token CE loss. + + Args: + logits: [B, T, V] float tensor. + labels: [B, T] int64 next-token labels; LABEL_IGNORE_INDEX positions are masked. + + Returns: + Scalar loss tensor. + """ + + def compute( + self, + logits: torch.Tensor, + labels: torch.Tensor, + ) -> torch.Tensor: + return F.cross_entropy( + logits[:, :-1].reshape(-1, logits.shape[-1]), + labels[:, :-1].reshape(-1), + ignore_index=LABEL_IGNORE_INDEX, + ) + + +class KLDivergenceLoss: + """Hinton KL-divergence with temperature. Only active (non-LABEL_IGNORE_INDEX) positions. + + Args: + student_logits: [B, T, V] float tensor, grad attached. + teacher_logits: [B, T, V] float tensor, no grad. + labels: [B, T] int64 next-token labels; LABEL_IGNORE_INDEX is masked. + """ + + def __init__(self, temperature: float = 1.0) -> None: + self.temperature = temperature + + def compute( + self, + student_logits: torch.Tensor, + teacher_logits: torch.Tensor, + labels: torch.Tensor, + ) -> torch.Tensor: + valid_mask = labels[:, :-1] != LABEL_IGNORE_INDEX + valid_s = student_logits[:, :-1][valid_mask] + valid_t = teacher_logits[:, :-1][valid_mask] + if valid_s.shape[0] == 0: + return torch.zeros(1, device=student_logits.device).squeeze() + return F.kl_div( + F.log_softmax(valid_s / self.temperature, dim=-1), + F.softmax(valid_t / self.temperature, dim=-1), + reduction="batchmean", + ) * (self.temperature * self.temperature) diff --git a/examples/qualcomm/oss_scripts/llama/train/trainer.py b/examples/qualcomm/oss_scripts/llama/train/trainer.py new file mode 100644 index 00000000000..ee6954aeb05 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/train/trainer.py @@ -0,0 +1,303 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import math +import re +import time +from abc import ABC, abstractmethod +from typing import Callable, Dict, List, Optional, Tuple + +import torch +from executorch.examples.qualcomm.oss_scripts.llama.train.config import TrainingArgs +from executorch.examples.qualcomm.oss_scripts.llama.train.loss import ( + CrossEntropyLoss, + KLDivergenceLoss, +) +from executorch.examples.qualcomm.oss_scripts.llama.train.utils import ( + build_param_groups, + get_warmup_cosine_lr, +) +from torch.utils.data import DataLoader +from tqdm import tqdm + + +class BaseTrainer(ABC): + """Training loop skeleton: epoch/batch iteration, grad accumulation, clipping, optimizer step. + + Subclasses implement train_step, eval_step, and configure_optimizers. + """ + + def __init__( + self, + model: torch.nn.Module, + args: TrainingArgs, + train_loader: DataLoader, + val_loader: Optional[DataLoader] = None, + ) -> None: + self.model = model + self.args = args + self.train_loader = train_loader + self.val_loader = val_loader + self.device = torch.device("cpu") # TODO: validate GPU compatibility + self._optimizer: Optional[torch.optim.Optimizer] = None + self._scheduler = None + + @abstractmethod + def train_step( + self, + model_inputs: Dict[str, torch.Tensor], + ) -> Tuple[torch.Tensor, Dict[str, float]]: + """Forward + loss. Returns (loss, metrics). Base calls .backward() — do NOT call it here.""" + + @abstractmethod + def eval_step( + self, + model_inputs: Dict[str, torch.Tensor], + ) -> Dict[str, float]: + """Validation forward for one batch. No gradients.""" + + @abstractmethod + def configure_optimizers( + self, warmup_steps: int, total_steps: int + ) -> Tuple[torch.optim.Optimizer, object]: + """Build (optimizer, lr_scheduler). Called after model.to(device).""" + + def train(self) -> torch.nn.Module: + self.model.to(self.device) + + num_batches = len(self.train_loader) + optimizer_steps_per_epoch = math.ceil(num_batches / self.args.grad_accum_steps) + total_steps = optimizer_steps_per_epoch * self.args.epochs + warmup_steps = max(2, int(total_steps * self.args.warmup_ratio)) + self._optimizer, self._scheduler = self.configure_optimizers( + warmup_steps, total_steps + ) + + if self._optimizer is None: + logging.info( + f"{type(self).__name__}: no trainable parameters, training loop skipped." + ) + self.model.to("cpu") + return self.model + + logging.info( + f"{type(self).__name__}: device={self.device} " + f"train_batches={num_batches} epochs={self.args.epochs} " + f"steps_per_epoch={optimizer_steps_per_epoch} warmup={warmup_steps}" + ) + + for epoch in tqdm(range(self.args.epochs)): + total_loss = 0.0 + t0 = time.time() + + for batch_idx, model_inputs in enumerate(self.train_loader): + with torch.enable_grad(): + loss, metrics = self.train_step(model_inputs) + (loss / self.args.grad_accum_steps).backward() + + is_last_batch = batch_idx == num_batches - 1 + if (batch_idx + 1) % self.args.grad_accum_steps == 0 or is_last_batch: + if self.args.max_grad_norm > 0: + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), self.args.max_grad_norm + ) + self._optimizer.step() + self._scheduler.step() + self._optimizer.zero_grad() + + lr = self._scheduler.get_last_lr()[0] + logging.info( + f"Epoch {epoch}, batch {batch_idx}/{num_batches} " + + " ".join(f"{k}={v:.4f}" for k, v in metrics.items()) + + f" lr={lr:.2e}" + ) + total_loss += loss.item() + + val_info = "" + if self.val_loader is not None: + val_metrics = self._run_validation() + val_info = " " + " ".join( + f"val_{k}={v:.4f}" for k, v in val_metrics.items() + ) + logging.info( + f"Epoch {epoch}: avg_loss={total_loss / max(num_batches, 1):.4f}" + f"{val_info} elapsed={time.time() - t0:.1f}s" + ) + + self.model.to("cpu") + return self.model + + def _to_device( + self, model_inputs: Dict[str, torch.Tensor] + ) -> Dict[str, torch.Tensor]: + return { + k: v.to(self.device) if isinstance(v, torch.Tensor) else v + for k, v in model_inputs.items() + } + + @torch.no_grad() + def _run_validation(self) -> Dict[str, float]: + accum: Dict[str, float] = {} + n = 0 + for model_inputs in self.val_loader: + for k, v in self.eval_step(model_inputs).items(): + accum[k] = accum.get(k, 0.0) + v + n += 1 + return {k: v / n for k, v in accum.items()} if n > 0 else {} + + +class Trainer(BaseTrainer): + """CE-only trainer (SFT / PTQ-warmup fine-tuning). + + forward_fn and loss_fn are injected so this class is unaware of attn_masks or KV caches. + """ + + def __init__( + self, + model: torch.nn.Module, + args: TrainingArgs, + train_loader: DataLoader, + forward_fn: Callable, + val_loader: Optional[DataLoader] = None, + frozen_param_patterns: Optional[List[str]] = None, + ) -> None: + super().__init__(model, args, train_loader, val_loader) + self.forward_fn = forward_fn + self.loss_fn = CrossEntropyLoss() + self._frozen_param_patterns: List[str] = frozen_param_patterns or [] + + def configure_optimizers( + self, warmup_steps: int, total_steps: int + ) -> Tuple[torch.optim.Optimizer, object]: + self._freeze_params() + params = build_param_groups(self.model, self.args.lr_config) + if not params: + logging.info( + f"{type(self).__name__}: no trainable parameters, skipping optimizer setup." + ) + return None, None + optimizer = torch.optim.AdamW( + params, + lr=self.args.lr, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=0.01, + ) + scheduler = get_warmup_cosine_lr( + optimizer, + warmup_min_lr=1e-8, + warmup_max_lr=self.args.lr, + warmup_num_steps=max(2, warmup_steps), + total_steps=total_steps, + ) + logging.info( + f"{type(self).__name__}: AdamW lr={self.args.lr:.2e} " + f"weight_decay=0.01 loss={type(self.loss_fn).__name__}" + ) + return optimizer, scheduler + + def train_step( + self, + model_inputs: Dict[str, torch.Tensor], + ) -> Tuple[torch.Tensor, Dict[str, float]]: + model_inputs = self._to_device(model_inputs) + gt_labels = model_inputs["labels"] + logits = self.forward_fn(self.model, model_inputs) + loss = self.loss_fn.compute(logits, gt_labels) + return loss, {"ce": loss.item()} + + @torch.no_grad() + def eval_step(self, model_inputs: Dict[str, torch.Tensor]) -> Dict[str, float]: + model_inputs = self._to_device(model_inputs) + gt_labels = model_inputs["labels"] + logits = self.forward_fn(self.model, model_inputs) + loss = self.loss_fn.compute(logits, gt_labels) + return {"ce": loss.item()} + + def _freeze_params(self) -> None: + frozen = 0 + for name, param in self.model.named_parameters(): + if any(re.search(p, name) for p in self._frozen_param_patterns): + param.requires_grad_(False) + frozen += 1 + logging.info(f"QAT: frozen param {name}") + logging.info(f"QAT: frozen {frozen} params") + + +class KDTrainer(Trainer): + """QAT knowledge-distillation trainer. Extends Trainer with an FP32 teacher and KD loss.""" + + def __init__( + self, + model: torch.nn.Module, + teacher: torch.nn.Module, + args: TrainingArgs, + train_loader: DataLoader, + forward_fn: Callable, + val_loader: Optional[DataLoader] = None, + frozen_param_patterns: Optional[List[str]] = None, + ) -> None: + super().__init__( + model=model, + args=args, + train_loader=train_loader, + forward_fn=forward_fn, + val_loader=val_loader, + frozen_param_patterns=frozen_param_patterns, + ) + self.teacher = teacher + + self._kd_loss = KLDivergenceLoss(temperature=args.temperature) + self._alpha = args.alpha + + def configure_optimizers( + self, warmup_steps: int, total_steps: int + ) -> Tuple[torch.optim.Optimizer, object]: + optimizer, scheduler = super().configure_optimizers(warmup_steps, total_steps) + if optimizer is None: + return None, None + logging.info( + f"KDTrainer: alpha={self._alpha:.3f} (KLDivergenceLoss) " + f"(1-alpha)={1.0 - self._alpha:.3f} (CrossEntropyLoss)" + ) + return optimizer, scheduler + + def train(self) -> torch.nn.Module: + self.teacher.to(self.device) + self.teacher.eval() + + try: + return super().train() + finally: + self.teacher.to("cpu") + + def train_step( + self, + model_inputs: Dict[str, torch.Tensor], + ) -> Tuple[torch.Tensor, Dict[str, float]]: + model_inputs = self._to_device(model_inputs) + gt_labels = model_inputs["labels"] + + with torch.no_grad(): + teacher_logits = self.forward_fn(self.teacher, model_inputs) + + student_logits = self.forward_fn(self.model, model_inputs) + kd = self._kd_loss.compute(student_logits, teacher_logits, gt_labels) + ce = self.loss_fn.compute(student_logits, gt_labels) + loss = self._alpha * kd + (1.0 - self._alpha) * ce + return loss, {"kd": kd.item(), "ce": ce.item()} + + @torch.no_grad() + def eval_step(self, model_inputs: Dict[str, torch.Tensor]) -> Dict[str, float]: + model_inputs = self._to_device(model_inputs) + gt_labels = model_inputs["labels"] + + teacher_logits = self.forward_fn(self.teacher, model_inputs) + student_logits = self.forward_fn(self.model, model_inputs) + kd = self._kd_loss.compute(student_logits, teacher_logits, gt_labels) + ce = self.loss_fn.compute(student_logits, gt_labels) + return {"kd": kd.item(), "ce": ce.item()} diff --git a/examples/qualcomm/oss_scripts/llama/train/utils.py b/examples/qualcomm/oss_scripts/llama/train/utils.py new file mode 100644 index 00000000000..8bae9f0af0a --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/train/utils.py @@ -0,0 +1,96 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import fnmatch +import logging +import math +from typing import Dict, List, Optional + +import torch +import yaml + + +def get_warmup_cosine_lr( + optimizer: torch.optim.Optimizer, + warmup_min_lr: float, + warmup_max_lr: float, + warmup_num_steps: int, + total_steps: int, +) -> torch.optim.lr_scheduler.LambdaLR: + """Linear warm-up from warmup_min_lr to warmup_max_lr, then cosine decay to 0. + + Matches HuggingFace's `lr_scheduler_type: cosine` (single half-cycle, no min-lr floor). + """ + min_ratio = warmup_min_lr / warmup_max_lr + warmup_slope = (1.0 - min_ratio) / max(1, warmup_num_steps) + decay_steps = max(1, total_steps - warmup_num_steps) + + def lr_lambda(step: int) -> float: + if step < warmup_num_steps: + return min_ratio + step * warmup_slope + progress = min(1.0, (step - warmup_num_steps) / decay_steps) + return 0.5 * (1.0 + math.cos(math.pi * progress)) + + return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) + + +def _numeric_str_to_float(value): + # PyYAML parses bare scientific notation (e.g. "4e-6") as a string, not a + # float, unless it has a decimal point (e.g. "4.0e-6"). Convert explicitly so + # an lr_config entry like `lr: 4e-6` doesn't silently reach the optimizer as + # a str (AdamW accepts a str lr at construction time but crashes on .step()). + if isinstance(value, str): + try: + return float(value) + except ValueError: + return value + return value + + +def build_param_groups( + model: torch.nn.Module, lr_config_path: Optional[str] +) -> List[Dict]: + """Split params into AdamW groups, giving each group its own kwargs. + + lr_config_path is a YAML file mapping glob patterns to per-group + optimizer kwargs (see train/config/lr_config.yaml), matched against + `model.named_parameters()` names. + """ + trainable = [(n, p) for n, p in model.named_parameters() if p.requires_grad] + + pattern_cfg = {} + if lr_config_path: + with open(lr_config_path) as f: + pattern_cfg = yaml.safe_load(f) or {} + + used = set() + groups = [] + for pattern, overrides in pattern_cfg.items(): + matched = [ + (n, p) + for n, p in trainable + if n not in used and fnmatch.fnmatchcase(n, pattern) + ] + if not matched: + continue + used.update(n for n, _ in matched) + group = {k: _numeric_str_to_float(v) for k, v in overrides.items()} + if "betas" in group: + group["betas"] = tuple(group["betas"]) + group["params"] = [p for _, p in matched] + groups.append(group) + logging.info( + f"lr_config: pattern '{pattern}' matched {len(matched)} param(s) -> {overrides}" + ) + + default_params = [p for n, p in trainable if n not in used] + if default_params: + groups.append({"params": default_params}) + logging.info( + f"lr_config: {len(default_params)} param(s) using default optimizer settings" + ) + + return groups diff --git a/examples/qualcomm/oss_scripts/llama/wrappers/attention_sink_wrappers.py b/examples/qualcomm/oss_scripts/llama/wrappers/attention_sink_wrappers.py index da3a165277f..e81d2a7d55e 100644 --- a/examples/qualcomm/oss_scripts/llama/wrappers/attention_sink_wrappers.py +++ b/examples/qualcomm/oss_scripts/llama/wrappers/attention_sink_wrappers.py @@ -339,7 +339,7 @@ def quantize(self, request: Request): ) num_data = 200 data_generator = partial( - request.method_data[ATTENTION_SINK_EVICTOR].calibration_data.datasets, + request.method_data[ATTENTION_SINK_EVICTOR].quantization_data.calib_loader, kcache_shape=self.kv_cache_shape["k"], vcache_shape=self.kv_cache_shape["v"], ) @@ -415,8 +415,8 @@ def quantize( kv_quant_attrs=kv_quant_attrs, ), ), - calibration_data=Request.CalibrationData( - datasets=partial( + quantization_data=Request.QuantizationData( + calib_loader=partial( _calibration_data_generator, kv_quant_attrs=kv_quant_attrs ), ), diff --git a/examples/qualcomm/oss_scripts/llama/wrappers/base_component.py b/examples/qualcomm/oss_scripts/llama/wrappers/base_component.py index 504a3d3e9e4..1fc08163912 100644 --- a/examples/qualcomm/oss_scripts/llama/wrappers/base_component.py +++ b/examples/qualcomm/oss_scripts/llama/wrappers/base_component.py @@ -165,17 +165,19 @@ def process(self, request: Any): @dataclass class Request: @dataclass - class CalibrationData: - datasets: Optional[DataLoader] = None + class QuantizationData: + calib_loader: Optional[DataLoader] = None intermediate_outputs: Optional[DataLoader] = None qdq_intermediate_outputs: Optional[DataLoader] = None + train_loader: Optional[DataLoader] = None + val_loader: Optional[DataLoader] = None @dataclass class Data: compile_spec: List[CompileSpec] = None pte_filename: str = None custom_annotation: Any = () - calibration_data: Request.CalibrationData = None + quantization_data: Request.QuantizationData = None tokenizer: callable = None skip_quantize: bool = False backend: QnnExecuTorchBackendType = QnnExecuTorchBackendType.kHtpBackend diff --git a/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py b/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py index 9a1ad50f854..382cce4ecf1 100644 --- a/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py +++ b/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py @@ -85,11 +85,15 @@ LlamaModel, ModelArgs, ) -from executorch.examples.qualcomm.oss_scripts.llama.quantize import PTQStrategy +from executorch.examples.qualcomm.oss_scripts.llama.quantize import ( + PTQStrategy, + QATStrategy, +) from executorch.examples.qualcomm.oss_scripts.llama.static_llm_quant_recipe import ( StaticLLMQuantRecipe, ) from executorch.examples.qualcomm.oss_scripts.llama.tokenizer import TokenizerWrapper +from executorch.examples.qualcomm.oss_scripts.llama.train.config import TrainingArgs from executorch.examples.qualcomm.oss_scripts.llama.wrappers.base_component import ( Component, get_model_specific_kwargs, @@ -107,8 +111,12 @@ from executorch.extension.llm.export.builder import DType from torch.utils.data import DataLoader from torchao.prototype.spinquant import apply_spinquant -from torchao.quantization.pt2e import MinMaxObserver -from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e +from torchao.quantization.pt2e import MinMaxObserver, move_exported_model_to_train +from torchao.quantization.pt2e.quantize_pt2e import ( + convert_pt2e, + prepare_pt2e, + prepare_qat_pt2e, +) from transformers import AutoModel, AutoModelForSpeechSeq2Seq @@ -164,10 +172,13 @@ def __init__( self.pass_manager_cls.get_passes_dependency_for_capture_program() ) self.meta = {} + recipe_cls = ( + self.config.qat_recipe + if control_args.qat and self.config.qat_recipe + else self.config.quant_recipe + ) self.quant_recipe: StaticLLMQuantRecipe = ( - self.config.quant_recipe(mode == Mode.CALIBRATE) - if self.config.quant_recipe - else None + recipe_cls(mode == Mode.CALIBRATE) if recipe_cls else None ) # For multimodal embedding @@ -579,6 +590,7 @@ def quantize(self, request: Request): # noqa: C901 soc_model=data.soc_model, ) + use_qat = self.control_args.qat and self.mode == Mode.CALIBRATE with torch.no_grad(): graph_module = None self.decoder = torch.export.export( @@ -587,6 +599,7 @@ def quantize(self, request: Request): # noqa: C901 if ( self.mode == Mode.CALIBRATE and self.control_args.quant_recipe_suggestion + or use_qat ): graph_module = copy.deepcopy(self.decoder) if self.apply_embedding: @@ -614,7 +627,11 @@ def quantize(self, request: Request): # noqa: C901 event_name="export_tasks", ) - self.decoder = prepare_pt2e(self.decoder, quantizer) + if use_qat: + self.decoder = prepare_qat_pt2e(self.decoder, quantizer) + move_exported_model_to_train(self.decoder) + else: + self.decoder = prepare_pt2e(self.decoder, quantizer) if self.apply_embedding: self.tok_embedding = prepare_pt2e( self.tok_embedding, tok_embedding_quantizer @@ -624,19 +641,47 @@ def quantize(self, request: Request): # noqa: C901 calibration_dataloaders = { AUDIO_ENCODER: request.method_data[ AUDIO_ENCODER - ].calibration_data.intermediate_outputs, + ].quantization_data.intermediate_outputs, VISION_ENCODER: request.method_data[ VISION_ENCODER - ].calibration_data.intermediate_outputs, - TEXT_DECODER: data.calibration_data.datasets, + ].quantization_data.intermediate_outputs, + TEXT_DECODER: data.quantization_data.calib_loader, } - PTQStrategy( - inference=self._decoder_inference, - module=self.decoder, - seq_mse_candidates=self.config.seq_mse_candidates, - tok_embedding=self.tok_embedding, - ).quantize(calib_loader=calibration_dataloaders) - logging.info("Calibration complete for prepare_pt2e") + + if use_qat: + training_args = TrainingArgs.from_yaml( + self.control_args.train_config + ) + training_args.lr_config = self.control_args.lr_config + frozen = ( + [".*"] + if self.control_args.freeze_all_params + # freeze_all_params: CI-only flag to verify QAT vs PTQ accuracy difference + # by disabling weight updates and only updating scale/zero_point. + else getattr(self.quant_recipe, "frozen_param_patterns", None) + ) + QATStrategy( + inference=self._decoder_inference, + module=self.decoder, + tok_embedding=self.tok_embedding, + seq_mse_candidates=self.config.seq_mse_candidates, + ).quantize( + calib_loader=calibration_dataloaders, + training_args=training_args, + teacher=graph_module, + train_loader=data.quantization_data.train_loader, + val_loader=data.quantization_data.val_loader, + frozen_param_patterns=frozen, + ) + logging.info("QAT training complete") + else: + PTQStrategy( + inference=self._decoder_inference, + module=self.decoder, + tok_embedding=self.tok_embedding, + seq_mse_candidates=self.config.seq_mse_candidates, + ).quantize(calib_loader=calibration_dataloaders) + logging.info("Calibration complete") else: # one dummy inference to remove affine observer # error happened in convert_pt2e @@ -655,7 +700,7 @@ def quantize(self, request: Request): # noqa: C901 self.quant_recipe.recipe, ) - # FP32 model used for quant-recipe-suggestion reference; release after use. + # FP32 model used as QAT teacher or quant-recipe-suggestion reference; release after use. del graph_module gc.collect() @@ -1215,7 +1260,7 @@ def quantize(self, request: Request): return request_data = request.method_data[self.modality] - calibration_datasets = request_data.calibration_data.datasets + calibration_datasets = request_data.quantization_data.calib_loader with torch.no_grad(): self.model = torch.export.export(self.model, self.example_input).module() @@ -1223,7 +1268,7 @@ def quantize(self, request: Request): if request_data.skip_quantize: logging.info(f"skipping encoder quantization for {self.modality}") intermediate_outputs = self._calibrate(self.model, calibration_datasets) - request_data.calibration_data.intermediate_outputs = ( + request_data.quantization_data.intermediate_outputs = ( intermediate_outputs ) return @@ -1236,7 +1281,7 @@ def quantize(self, request: Request): # start calibration intermediate_outputs = self._calibrate(self.model, calibration_datasets) - request_data.calibration_data.intermediate_outputs = intermediate_outputs + request_data.quantization_data.intermediate_outputs = intermediate_outputs self.model = convert_pt2e(self.model) @@ -1245,7 +1290,7 @@ def quantize(self, request: Request): qdq_intermediate_outputs = self._calibrate( self.model, calibration_datasets ) - request_data.calibration_data.qdq_intermediate_outputs = ( + request_data.quantization_data.qdq_intermediate_outputs = ( qdq_intermediate_outputs ) @@ -1324,14 +1369,23 @@ def quantize( tokenizer_wrapper=tokenizer_wrapper, attn_mask=self.text_decoder.calibration_prefill.attn_mask, ) - calibration_data = dataset_builder.build_calib_dataloaders() + if self.control_args.qat: + calib_loader, train_loader, val_loader = ( + dataset_builder.build_qat_dataloaders() + ) + else: + calib_loader = dataset_builder.build_calib_dataloaders() + train_loader = dict.fromkeys(calib_loader) + val_loader = dict.fromkeys(calib_loader) quantize_request = Request( inspect.currentframe().f_code.co_name, { m: Request.Data( - calibration_data=Request.CalibrationData( - datasets=calibration_data[m] + quantization_data=Request.QuantizationData( + calib_loader=calib_loader[m], + train_loader=train_loader[m], + val_loader=val_loader[m], ), skip_quantize=skip_quantize.get(m, False), tokenizer=tokenizer_wrapper.tokenizer,