diff --git a/docs/source/rfc_llm_config_calibration_datasets.md b/docs/source/rfc_llm_config_calibration_datasets.md new file mode 100644 index 00000000000..f2e7daebf5c --- /dev/null +++ b/docs/source/rfc_llm_config_calibration_datasets.md @@ -0,0 +1,262 @@ +# RFC: Hugging Face calibration dataset in `LlmConfig` + +**Status:** RFC + +**Author:** Mahesh Madhavan + +**Last Update:** 2026-10-01 + +## Summary + +Add `quantization.calibration_dataset` to `LlmConfig` for Hugging Face datasets or +local files. A shared loader returns text/chat rows with optional image or audio +values. + +## Motivation + +`QuantizationConfig` currently exposes a single prompt and `lm_eval` task fields: + +```python +calibration_tasks: Optional[List[str]] = None +calibration_limit: Optional[int] = None +calibration_seq_length: Optional[int] = None +calibration_data: str = "Once upon a time" +``` + +These fields cannot describe an HF dataset or local dataset files. + +## Scope + +- Add a YAML- and OmegaConf-compatible dataset source using Huggingface `datasets.load_dataset()`. +- Describe text/chat, image+text, and audio+text inputs through named columns. +- Preserve existing prompt and `lm_eval` fields and behavior when the new field + is absent. + +The multimodal source contract and examples are included here. + +## Proposed configuration + +Add a YAML- and OmegaConf-compatible `HFDatasetConfig` to +[llm_config.py](../../extension/llm/export/config/llm_config.py), exposed as +`llm_config.quantization.calibration_dataset`. + +```diff +-from typing import ClassVar, List, Optional ++from typing import Any, ClassVar, Dict, List, Optional + ++class DatasetInputFormat(str, Enum): ++ text = "text" ++ messages = "messages" ++ ++ ++@dataclass ++class HFDatasetConfig: ++ # Forwarded to datasets.load_dataset(). ++ path: str ++ name: Optional[str] = None ++ data_dir: Optional[str] = None ++ data_files: Any = None ++ split: str = "train" ++ streaming: bool = False ++ load_dataset_kwargs: Dict[str, Any] = field(default_factory=dict) ++ ++ # Applied after loading. ++ limit: Optional[int] = 1 ++ input_format: DatasetInputFormat = DatasetInputFormat.text ++ input_column: str = "text" ++ image_column: Optional[str] = None ++ audio_column: Optional[str] = None + + @dataclass + class QuantizationConfig: + calibration_tasks: Optional[List[str]] = None + calibration_limit: Optional[int] = None + calibration_seq_length: Optional[int] = None + calibration_data: str = "Once upon a time" ++ calibration_dataset: Optional[HFDatasetConfig] = None +``` + +### Field summary + +| Field | Purpose | +|---|---| +| `path` | HF dataset ID or local builder such as `json` or `parquet`. | +| `name` | Optional dataset subset or builder configuration. | +| `data_dir` / `data_files` | Local or builder-specific data locations accepted by Hugging Face. | +| `split` | One dataset split to load. | +| `streaming` | Request a Hugging Face `IterableDataset`. | +| `load_dataset_kwargs` | Additional YAML-compatible arguments to `load_dataset()`. | +| `limit` | Maximum raw rows selected; defaults to 1. | +| `input_format` | Whether the selected input is plain text or messages. | +| `input_column` | Column containing the text or messages. | +| `image_column` / `audio_column` | Optional column containing the media associated with the same row. | + +All values must be YAML-compatible. A supplied `limit` must be a positive integer. +`limit: null` selects the full split for a non-streaming dataset; streaming requires +a positive limit. Skipped rows count toward the limit. The limit does +not count token sequences or batches. + +## Examples + +### Plain text + +```yaml +quantization: + calibration_dataset: + path: Salesforce/wikitext + name: wikitext-2-raw-v1 + split: train + streaming: true + limit: 100 + input_format: text + input_column: text +``` + +Example row returned by Hugging Face: + +```json +{"text": "Eight FAB subtypes were proposed in 1976 ."} +``` + +For local files, use `path: json` or `path: parquet` with `data_files`. For chat, +use `input_format: messages` and select the column containing the conversations. + +### Image and text + +```yaml +quantization: + calibration_dataset: + path: json + data_files: /datasets/calibration.jsonl + split: train + limit: 100 + input_format: text + input_column: prompt + image_column: image +``` + +Example JSONL row: + +```json +{"prompt": "What is shown?", "image": "/datasets/images/example.jpg"} +``` + +### Audio and messages from Parquet + +```yaml +quantization: + calibration_dataset: + path: parquet + data_files: /datasets/audio_calibration.parquet + split: train + limit: 100 + input_format: messages + input_column: messages + audio_column: audio +``` + +Example record stored in Parquet (shown as JSON for readability): + +```json +{ + "messages": [ + {"role": "user", "content": "Transcribe this recording."}, + {"role": "assistant", "content": "The train leaves at six."} + ], + "audio": "/datasets/audio/example.wav" +} +``` + +This example associates one recording with the user turn. The model must +support that association, insert any required audio markers, apply chat template, +and process the recording using the model's audio requirements. The shared loader +passes the messages and audio value together without adding model-specific tokens. + +The Parquet file stores the nested messages and an audio file reference directly. + +### Audio and transcription directly from Hugging Face + +```yaml +quantization: + calibration_dataset: + path: google/fleurs + name: en_us + split: train + streaming: true + limit: 100 + input_format: text + input_column: transcription + audio_column: audio +``` + +FLEURS supplies `audio` and `transcription`, without a `messages` column. +Before preprocessing, a decoded row looks like this. + +```python +raw_row = { + "transcription": "The train leaves at six.", + "audio": { + "path": ".../example.wav", + "array": waveform, # Decoded audio samples. + "sampling_rate": 16000, + }, +} +``` + +An offline conversion can add messages while retaining the audio: + +```python +chat_row = { + "messages": [ + {"role": "user", "content": "Transcribe this recording."}, + {"role": "assistant", "content": raw_row["transcription"]}, + ], + "audio": raw_row["audio"], # Retain the audio with its sampling rate. +} +``` + +For the converted dataset, select `input_format: messages` and +`input_column: messages`. Audio can remain decoded or be saved as a WAV file at +its sampling rate, with the file path stored in `audio`. + +## Input contract + +The shared loader returns a native HF `Dataset` or `IterableDataset`, preserving +original column names: + +```python +def load_hf_dataset(config: HFDatasetConfig) -> Dataset | IterableDataset: + ... +``` + +Model integrations consume the rows as `Iterable[Mapping[str, Any]]`, with the +source configuration supplied separately: + +- For `text`, `input_column` selects a string. +- For `messages`, pass the value unchanged to the selected model's chat-template + adapter, which validates its supported message format. Shared code does not infer + or rename roles. +- Media columns select values from the same row. Preserve the representation + returned by HF. + +The shared loader performs no additional media normalization or model preprocessing. + +## Compatibility + +- Without `calibration_dataset`, preserve existing prompt and `lm_eval` behavior. +- With it, use the dataset as the sole source; reject other explicitly supplied + calibration sources. +- HF `limit` controls rows independently of `lm_eval` task limits. + +## Implementation direction + +1. Add the config and source-selection validation. +2. Add a shared LLM loader that lazily imports `datasets`, loads one split, and + applies the row limit. +3. Perform any optional dataset normalization offline. + +## Decision requested + +1. Add `quantization.calibration_dataset: Optional[HFDatasetConfig]` without + changing existing calibration fields. +2. Use native HF datasets and dict rows as the shared source interface. diff --git a/docs/source/rfc_llm_config_lora.md b/docs/source/rfc_llm_config_lora.md new file mode 100644 index 00000000000..a174b6e057f --- /dev/null +++ b/docs/source/rfc_llm_config_lora.md @@ -0,0 +1,221 @@ +# RFC: LoRA methods and calibration in `LlmConfig` + +**Status:** RFC; implementation pending. + +**Author:** Mahesh Madhavan + +**Last Update:** 2026-10-01 + +**Prerequisite:** [Hugging Face calibration dataset support](rfc_llm_config_calibration_datasets.md). + +## Summary + +Let users choose LoRA adapters, quantization settings, and calibration data for +each exported method through `LlmConfig`. + +A method is a named entry point for calling the exported model. It can use the +base model alone or include a LoRA adapter: a small set of trained weights that +changes the model's behavior. + +## What changes + +Today, `MethodConfig.lora_config` selects an adapter. This proposal adds: + +- `lora_configs`: a list of adapters for each method, with optional names. +- `quantization.lora_quantize`: optional quantization and dataset settings for a + method, selected by `method_name`. +- `HFDatasetConfig`: the same dataset config used by the calibration dataset RFC. +- `backend.qnn.lora`: QNN settings for passing adapter inputs and keeping + quantization values fixed. + +V1 allows several methods, with one lora adapter, and supports only +`freeze_all` quantization_strategy. More adapters per method and other strategies are future work. + +## Proposed configuration + +Changes to `extension/llm/export/config/llm_config.py`, after the prerequisite RFC: + +```diff + @dataclass + class LoraConfig: + # Existing fields remain unchanged. ++ name: Optional[str] = None + + @dataclass + class MethodConfig: + method_name: str + lora_config: Optional[LoraConfig] = None + export_seq_len: Optional[int] = None ++ lora_configs: List[LoraConfig] = field(default_factory=list) + ++@dataclass ++class LoraQuantizationConfig: ++ method_name: str ++ lora_pt2e_quantize: Optional[Pt2eQuantize] = None ++ calibration_dataset: Optional[HFDatasetConfig] = None + + @dataclass + class QuantizationConfig: + # Includes calibration_dataset from the prerequisite RFC. ++ lora_quantize: List[LoraQuantizationConfig] = field(default_factory=list) + ++class QNNLoraQuantizationStrategy(str, Enum): ++ freeze_all = "freeze_all" + ++@dataclass ++class QNNLoraConfig: ++ weight_as_input: bool = True ++ scale_as_input: bool = True ++ quantization_strategy: QNNLoraQuantizationStrategy = ( ++ QNNLoraQuantizationStrategy.freeze_all ++ ) + + @dataclass + class QNNConfig: + # Existing fields remain unchanged. ++ lora: QNNLoraConfig = field(default_factory=QNNLoraConfig) +``` + +- Add `name` after the existing fields so existing calls to `LoraConfig(...)` + still work. +- Accept the old `lora_config` field for compatibility; use `lora_configs` in new + configs. Reject a method that sets both an adapter and a nonempty adapter list. +- If neither field supplies an adapter, the method uses only the base model. + +An adapter's `adapter_checkpoint` points to its trained weights. +`adapter_config` points to a file describing its rank, scaling, and target layers. + +## Choosing settings for each method + +| Setting | Behavior | +|---|---| +| `LoraConfig.name` | Optional adapter name. If omitted, use `method_name` for the method's single adapter. | +| `multimethod.methods[].lora_configs` | LoRA weights and config for this method. | +| `lora_quantize[].method_name` | Names the existing method whose settings to change. That method must have an adapter. | +| `lora_pt2e_quantize` | Quantization settings for the adapter. If omitted or `None`, use base model settings. | +| `calibration_dataset` | Dataset for this method. If omitted or `None`, use base model `quantization.calibration_dataset`. | + +A method without a `lora_quantize` entry uses the shared settings. + +## QNN settings and `freeze_all` + +These settings live under `backend.qnn.lora`: + +| Setting | Proposed meaning | +|---|---| +| `weight_as_input` | If true, pass adapter weights when calling the model. If false, include them in the exported method. | +| `scale_as_input` | Pass the adapter's scaling factor as an input. This controls how much the adapter contributes to the result. | +| `quantization_strategy` | `freeze_all` keeps both base and adapter quantization values fixed during export. | + + + +## Example + +This example uses the datasets to prepare quantization values before `freeze_all` +export. It requires exporter support; model and export-shape settings are omitted. + +```yaml +backend: + qnn: + enabled: true + lora: + weight_as_input: true + scale_as_input: true + quantization_strategy: freeze_all + +quantization: + pt2e_quantize: qnn_16a4w + calibration_dataset: + path: Salesforce/wikitext + name: wikitext-2-raw-v1 + split: train + limit: 100 + input_format: text + input_column: text + + lora_quantize: + - method_name: function_usecase + lora_pt2e_quantize: qnn_16a16w + calibration_dataset: + path: json + data_files: /datasets/function/calibration.jsonl + split: train + limit: 100 + input_format: messages + input_column: messages + - method_name: elementary_usecase + lora_pt2e_quantize: qnn_16a16w + # Inherits base calibration_dataset. + +multimethod: + methods: + - method_name: base_usecase + - method_name: function_usecase + lora_configs: + - name: function + adapter_checkpoint: /adapters/function/adapter_model.safetensors + adapter_config: /adapters/function/adapter_config.json + - method_name: elementary_usecase + lora_configs: + - name: elementary + adapter_checkpoint: /adapters/elementary/adapter_model.safetensors + adapter_config: /adapters/elementary/adapter_config.json +``` + +| Method | Adapter | Calibration dataset | +|---|---|---| +| `base_usecase` | None | Shared WikiText dataset | +| `function_usecase` | `function` | Method's chat dataset in JSONL | +| `elementary_usecase` | `elementary` | Shared WikiText dataset | + +Using the same dataset does not mean sharing collected calibration statistics or +combining adapter weights. Prepare values for each method's model and adapter. + +## Checks and compatibility + +- Method names must be unique. Allow at most one `lora_quantize` entry per method; + reject entries for missing methods or methods without adapters. +- Reject multiple adapters per method and strategies other than `freeze_all` in + v1. +- Keep existing Python/YAML configs working. Reject combining + `base.lora_config` with `multimethod.methods`. +- Keep `adapter_quant` behavior; reject conflicts with shared or per-method + quantization settings. + + +## Future work + +### Additional quantization strategies + +- `freeze_base`: keep base quantization values fixed and recalibrate the adapter. +- `freeze_none`: recalibrate both base and adapter quantization values. + + +### Multiple adapters in one method + +The same list can describe multiple adapters in one method. Supporting their +combined execution is future work: + +```yaml +multimethod: + methods: + - method_name: combined_usecase + lora_configs: + - name: function + adapter_checkpoint: /adapters/function/adapter_model.safetensors + adapter_config: /adapters/function/adapter_config.json + - name: elementary + adapter_checkpoint: /adapters/elementary/adapter_model.safetensors + adapter_config: /adapters/elementary/adapter_config.json +``` + +- Both adapters would run in one method, each with a unique name. +- Settings for `combined_usecase` would select one LoRA quantization setting and + one dataset for the whole method. + +## Decision requested + +1. Use `HFDatasetConfig` for each method's calibration dataset. +2. Add adapter lists and optional names while keeping the old `lora_config` field for now. +3. Add `backend.qnn.lora` with only `freeze_all` in v1, once the exporter can + prepare or load the required quantization values.