Update JointQ

Author: FKKimura

CHANGELOG.md

@@ -116,6 +116,18 @@
   - `chunking.py`: shared chunking strategies (`concat_chunk`, `concat_chunk_align`, `concat_rand`, `drop_head`, `drop_rand`) extracted as reusable helpers
 - Added `calibration_dataset` parameter to `AutoBitQuantizer` to specify the calibration data source (`onecomp/quantizer/autobit/_autobit.py`)
 
+### JointQ:
+
+- Added `incremental_lambda` regularization mode (`lambda_mode="incremental_lambda"`): for each layer, tries increasing lambda values from `lambda_list` with warm start, accepting candidates that improve weight error without substantially degrading output error. Stops at the first rejection. Controlled by `lambda_list`, `incremental_eps_y`, `incremental_eps_w` parameters
+- Added `incremental_initial_skip_ew_threshold` to skip an unstable initial `lambda=0.0` candidate when its relative weight error is excessively large
+- Added `accepted_lambda` field to `JointQResult` to record the per-layer lambda chosen in incremental mode
+- Added `execute_post_processing` override to log accepted lambda statistics (mean, median, min, max, per-value counts) after all layers are quantized
+- Added `regularization_mode` parameter: `"identity"` (standard Tikhonov λI) or `"diagonal"` (default, importance-aware λ·diag(a) where a_i scales with activation magnitude). Diagonal mode reduces over-regularization of less important columns. Only supported with `fixed_lambda` mode
+- Added `regularization_gamma` parameter (default 0.5): exponent for diagonal weights in `"diagonal"` mode; smaller values reduce the spread between weak and strong columns
+- Added initialization strategy control: `enable_clip_optimize`, `enable_clip_optimize_ep`, `enable_gptq` parameters to `JointQ` class
+- Added `gptq` attribute (`GPTQ` instance) to `JointQ` class for customizing GPTQ parameters (`blocksize`, `percdamp`, `mse`, `q_grid`, `q_norm`). Default GPTQ is auto-created from `bits`/`group_size`/`symmetric`
+- Replaced JointQ internal GPTQ module (`jointq/core/gptq.py`) with OneComp GPTQ (`onecomp.quantizer.gptq.GPTQ`); GPTQ initial solution is now generated via the shared GPTQ implementation
+
 ### Breaking Changes
 
 - **`AutoBitQuantizer.enable_fused_groups` now defaults to `True`** (`onecomp/quantizer/autobit/_autobit.py`)
@@ -126,12 +138,21 @@
 - Quantisation levels unified to unsigned `[0, 2^b − 1]` (symmetric uses centred zero point); rounding order changed from `round(x/s + z)` to `round(x/s) + z`. Outputs are not bit-exact with prior RTN versions
 - Changed `prepare_rotated_model` defaults: `rotation_mode` `"random"` → `"random_hadamard"`, `num_calibration_samples` `128` → `512`
 - Introduced `CalibrationConfig` dataclass; see CalibrationConfig section above for migration details
+- `JointQ` class: removed `batch_size` parameter (use `onecomp.quantizer.jointq.core.quantize()` directly if batch processing is needed)
+- `JointQ`: GPTQ initial solution is now generated via OneComp GPTQ instead of the internal implementation. Quantization results are not bit-exact with prior versions (quality is equivalent or improved)
+- **`JointQ` regularization defaults changed** (`onecomp/quantizer/jointq/_jointq.py`)
+  - `regularization_lambda`: `0.2` → `0.1`
+  - `regularization_mode`: `"identity"` → `"diagonal"`
+  - Quantization results are not bit-exact with prior versions.
+  - **Migration:** to reproduce the previous behavior, pass
+    `JointQ(..., regularization_lambda=0.2, regularization_mode="identity")` explicitly.
 
 ### Bug Fix
 
 - Fixed `model_config.py`: `load_model()` VLM fallback did not trigger for models raising `"Unrecognized configuration class"` (e.g. Cohere2VisionForConditionalGeneration). Added the error pattern to `_vlm_hints`
 - Fixed `gptq/_gptq.py`: Cholesky decomposition in `run_gptq` could fail with `LinAlgError` on ill-conditioned Hessians (observed on large VLMs at deeper layers). Extracted `_compute_inverse_hessian()` with progressive damping fallback (up to 5 retries, 10x damping increase per retry). No impact on normal operation
 - Fixed `TypeError` in `QuantLinear.forward` when `S_qk` scaling was applied to MLP layers (`onecomp/pre_process/quant_models.py`)
+- Fixed `JointQ` `group_size=None` (per-channel quantization) raising `TypeError`
 - Fixed wrong module grouping in `make_grouped_module` where GC-driven `id()` reuse caused attention projections (q/k/v) and MLP projections (gate/up) to be merged into the same group.  (`qep/_quantize_with_qep_arch.py`)
 - Fixed silent weight corruption in `GPTQLinear` when `qzero=0` was stored through the GPTQ v1 zero-point path (`onecomp/quantizer/gptq/gptq_layer.py`)
   - Root cause: AutoGPTQ v1 stores `raw_zero - 1`, so `qzero=0` becomes `-1`; without masking, its sign-extended bits corrupted neighboring packed slots
@@ -140,14 +161,6 @@
   - Added regression tests for per-slot pack corruption, packed/unpacked forward paths, the `gptq_v2` branch, and the `from_saved_state` path for GPTQ v1 tensors
   - **NOTE**: If you have GPTQ models quantized with previous versions, please re-quantize them with this release, as they may contain corrupted internal data.
 
-### Packaging
-
-- Bumped minimum `transformers` requirement from `>=5.3.0` to `>=5.5.0` (`pyproject.toml`)
-- Added `cu130` optional-dependency extra and the `pytorch-cu130` wheel index (`https://download.pytorch.org/whl/cu130`) for CUDA 13 hosts (e.g. NVIDIA B200) (`pyproject.toml`)
-- Pinned the `vllm` extra to `vllm>=0.10` to prevent uv from falling back to legacy versions whose source build requires `CUDA_HOME` (`pyproject.toml`)
-- Added uv `conflicts` declarations between the `vllm` extra and the `cpu` / `cu118` / `cu121` / `cu124` / `cu126` / `cu128` extras: vLLM `>=0.20` requires `torch>=2.10`, which is only published for `cu130`. This forces `vllm` to be installed only with `--extra cu130` and prevents silent fallback to a `vllm` version incompatible with `transformers>=5` at runtime (`pyproject.toml`)
-- Restricted `tool.uv.environments` to `sys_platform == 'linux'` and `python_full_version >= '3.12', < '3.14'` to skip lock splits for unused Windows and out-of-range Python versions (`pyproject.toml`)
-
 ### Examples
 
 - Added `example/example_custom_calibration.py`: Demonstrates `CalibrationConfig` with a custom calibration dataset (Python code snippets in `example/data/python_calibration.txt`).  Quantizes TinyLlama with GPTQ 3-bit using both default C4 and custom Python-code calibration, then compares inference outputs across multiple prompts to show how calibration data choice affects quantization quality.
@@ -158,6 +171,7 @@
 
 ### Documentation
 
+- Updated `docs/algorithms/jointq.md`: added incremental lambda mode description, acceptance criteria, diagonal regularization mode, new parameters (`lambda_mode`, `lambda_list`, `incremental_eps_y`, `incremental_eps_w`, `incremental_initial_skip_ew_threshold`, `regularization_mode`, `regularization_gamma`), and usage examples
 - Added `docs/algorithms/lpcd.md`: LPCD overview, motivation, supported submodule groups, usage examples, QEP relationship, parameters, and current support
 - Added `docs/api/lpcd_config.md` and updated `mkdocs.yml` navigation to include LPCD in the Algorithms and API Reference sections
 - Updated LPCD references across docs: `docs/index.md`, `docs/algorithms/overview.md`, `docs/user-guide/basic-usage.md`, `docs/user-guide/configuration.md`, `docs/user-guide/examples.md`, and `docs/getting-started/quickstart.md`
@@ -177,14 +191,17 @@
 - Added `docs/api/quantizers/onebit.md` (OneBit API reference)
 - Updated `mkdocs.yml` nav: added AutoBit/JointQ algorithm pages, OneBit API page; renamed Post-Process nav title to include Block-wise PTQ
 - Added example script links to `docs/user-guide/pre-process.md`
+- Updated `docs/algorithms/jointq.md`: added initialization strategy parameters, GPTQ customization examples, removed `batch_size` from parameter table, added `group_size=None` per-channel documentation
 - Added a Troubleshooting section to `docs/user-guide/vllm-inference.md` describing how to bypass the unconditional DeepGEMM (FP8) kernel warmup for non-FP8 quantization (GPTQ / DBF / Mixed-GPTQ) by setting `VLLM_USE_DEEP_GEMM=0` and `VLLM_DEEP_GEMM_WARMUP=skip`
 
 ### Tests
 
+- Updated `test_jointq.py`: added `incremental_lambda` boundary parameters (`lambda_mode`, `lambda_list`, `incremental_eps_y`, `incremental_eps_w`, `incremental_initial_skip_ew_threshold`), abnormal parameter cases (invalid mode, empty list, negative values), GPU integration tests (`test_incremental_lambda_basic`, `test_incremental_lambda_single_step`), `_accept_candidate` unit tests covering all acceptance rules and edge cases; removed `batch_size` from boundary/abnormal parameter tests
 - Added `test_prepare_rotated_model.py`: validation, E2E pipeline, output threshold (80 combinations), save/load round-trip
 - Added `test_weight_quantizer.py`: RTN/GPTQ consistency, symmetric/asymmetric, group-wise, MSE, STE
 - Expanded `test_rtn.py`: MSE boundary/abnormal parameters
 - Added vLLM mixed group-size tests (`tests/vllm_plugins/gptq/test_mixed_gptq.py`, `tests/vllm_plugins/gptq/test_mixed_gptq_e2e.py`)
+- Updated `regression_quantize_helper.py`: updated `EXPECTED_MSE` baseline for OneComp GPTQ integration
 - Added LPCD tests (`tests/onecomp/lpcd/`, 25 cases)
   - `test_lpcd_config.py`: `LPCDConfig` default / custom values, dataclass field set, top-level `from onecomp import LPCDConfig` (CPU only)
   - `test_lpcd_metrics.py`: `make_lpcd_metrics()` dispatch on synthetic Llama / Qwen3 blocks for every `enable_*` flag combination, `NotImplementedError` for unsupported architectures, `LpcdMetricGroup.mark_as_ready` / `is_refineable` state transitions (CPU only, no weight download)
@@ -205,6 +222,11 @@
 
 ### Packaging
 
+- Bumped minimum `transformers` requirement from `>=5.3.0` to `>=5.5.0` (`pyproject.toml`)
+- Added `cu130` optional-dependency extra and the `pytorch-cu130` wheel index (`https://download.pytorch.org/whl/cu130`) for CUDA 13 hosts (e.g. NVIDIA B200) (`pyproject.toml`)
+- Pinned the `vllm` extra to `vllm>=0.10` to prevent uv from falling back to legacy versions whose source build requires `CUDA_HOME` (`pyproject.toml`)
+- Added uv `conflicts` declarations between the `vllm` extra and the `cpu` / `cu118` / `cu121` / `cu124` / `cu126` / `cu128` extras: vLLM `>=0.20` requires `torch>=2.10`, which is only published for `cu130`. This forces `vllm` to be installed only with `--extra cu130` and prevents silent fallback to a `vllm` version incompatible with `transformers>=5` at runtime (`pyproject.toml`)
+- Restricted `tool.uv.environments` to `sys_platform == 'linux'` and `python_full_version >= '3.12', < '3.14'` to skip lock splits for unused Windows and out-of-range Python versions (`pyproject.toml`)
 - Added `hydra` extra to `pyproject.toml` so `hydra-core` (used by `example/example_autobit.py` and the `model_validation/{gptq,qep_gptq,autobit,autobit_qep,jointq}/validate_*.py` scripts) installs in one step via `uv sync --extra <cuXXX> --extra hydra` or `pip install "onecomp[hydra]"`, instead of a separate `pip install hydra-core` after sync. Documented the new extra in `README.md` and the `model_validation/*/README.md` files. The `model_validation/gptq/` Phase 3 (vLLM inference) additionally requires the `vllm` extra (`uv sync --extra <cuXXX> --extra hydra --extra vllm` or `pip install "onecomp[hydra]" vllm`).
 
 ## [v1.0.2] 2026-03-31

docs/algorithms/jointq.md

@@ -26,35 +26,95 @@ with group index \(g\) for group-wise quantization.
 
 ### Initialization Strategies
 
-JointQ supports multiple initialization strategies for the local search:
+JointQ supports multiple initialization strategies for the local search.
+Each strategy can be independently enabled or disabled:
 
-1. **Clip-Optimize**: Finds optimal clipping range, then quantizes
-2. **Clip-Optimize with Error Propagation**: Adds GPTQ-style error propagation to initialization
-3. **GPTQ**: Uses GPTQ solution as the starting point for joint optimization
+1. **Clip-Optimize** (`enable_clip_optimize`): Finds optimal clipping range, then quantizes
+2. **Clip-Optimize with Error Propagation** (`enable_clip_optimize_ep`): Adds GPTQ-style error propagation to initialization
+3. **GPTQ** (`enable_gptq`): Uses OneComp GPTQ solution as the starting point for joint optimization
+
+By default, Clip-Optimize and GPTQ are enabled, while Clip-Optimize with Error
+Propagation is disabled. The best solution among all enabled strategies is
+selected row-by-row based on the reconstruction error.
 
 ### Regularization
 
 To prevent overfitting to calibration data, JointQ applies Tikhonov regularization:
 
 \[
-X^T X + n \lambda I
+X^T X + n \lambda R
 \]
 
-where \(\lambda\) controls the regularization strength (default: 0.2). This stabilizes
-the optimization, especially when the number of calibration samples is small relative
-to the model dimension.
+where \(\lambda\) controls the regularization strength and \(R\) is the
+regularization matrix controlled by `regularization_mode`:
+
+- **`"identity"`**: \(R = I\). Standard Tikhonov — all input
+  dimensions are protected equally.
+- **`"diagonal"`** (default): \(R = \mathrm{diag}(a)\) where
+  \(a_i = \bigl(\mathrm{diag}(X^TX)_i \;/\; \mathrm{mean}(\mathrm{diag}(X^TX))\bigr)^{\gamma}\).
+  Columns with larger activations receive stronger regularization while
+  less important columns are given more freedom, reducing over-protection.
+  Only supported with `lambda_mode="fixed_lambda"`.
+
+JointQ provides two lambda selection modes, selected by `lambda_mode`.
+
+#### Fixed lambda mode (default)
+
+Uses a single fixed \(\lambda\) for all layers (`lambda_mode="fixed_lambda"`).
+The default strength is `regularization_lambda=0.1`.
+
+#### Incremental lambda mode
+
+`lambda_mode="incremental_lambda"` tries increasing \(\lambda\) values from
+`lambda_list` for each layer, keeping the solution as long as it improves
+weight error without substantially degrading output error.
+
+For each layer, the algorithm:
+
+1. Quantizes with the first (smallest) \(\lambda\).
+   - If the first candidate uses `lambda=0.0` and its relative weight error
+     is extremely large, JointQ can skip that candidate and move to the next
+     lambda via `incremental_initial_skip_ew_threshold`.
+2. For each subsequent \(\lambda\), re-quantizes using the previous accepted
+   solution as a warm start.
+3. Accepts the candidate if the relative weight error \(E_w\) decreases
+   without the relative output error \(E_y\) worsening beyond tolerance.
+4. Stops at the first rejection and returns the last accepted solution.
+
+The acceptance criteria are:
+
+- Both \(E_w\) and \(E_y\) decreased → accept.
+- \(E_w\) increased → reject.
+- \(E_y\) worsened within `incremental_eps_y` and \(E_w\) improved by at
+  least `incremental_eps_w` → accept.
+- Otherwise → reject.
+
+Where:
+
+- \(E_w = \|W_q - W\|_F^2 \;/\; \|W\|_F^2\)
+- \(E_y = \|(W_q - W) X^T\|_F^2 \;/\; \|W X^T\|_F^2\)
 
 ## Parameters
 
-| Parameter               | Type    | Description                                         | Default  |
-|-------------------------|---------|-----------------------------------------------------|----------|
-| `bits`                  | `int`   | Quantization bit-width (1--4)                       | `4`      |
-| `symmetric`             | `bool`  | Symmetric quantization                              | `False`  |
-| `group_size`            | `int`   | Group size for group-wise quantization              | `128`    |
-| `batch_size`            | `int`   | Batch size for processing rows (None = all at once) | `None`   |
-| `regularization_lambda` | `float` | Tikhonov regularization strength                    | `0.2`    |
-| `actorder`              | `bool`  | Reorder columns by activation magnitude             | `False`  |
-| `device`                | `torch.device` | Device for computation                       | `None`   |
+| Parameter               | Type           | Description                                                | Default  |
+|-------------------------|----------------|------------------------------------------------------------|----------|
+| `bits`                  | `int`          | Quantization bit-width (1--4)                              | `4`      |
+| `symmetric`             | `bool`         | Symmetric quantization                                     | `False`  |
+| `group_size`            | `int` or `None`| Group size for group-wise quantization (None = per-channel) | `128`    |
+| `lambda_mode`           | `str`          | `"fixed_lambda"` or `"incremental_lambda"`                 | `"fixed_lambda"` |
+| `regularization_lambda` | `float` or `None` | Tikhonov regularization strength (fixed mode)          | `0.1`    |
+| `regularization_mode`  | `str`              | `"identity"` (λI) or `"diagonal"` (λ·diag(a), fixed mode only) | `"diagonal"` |
+| `regularization_gamma` | `float`            | Exponent for diagonal weights (`"diagonal"` mode)      | `0.5`    |
+| `lambda_list`           | `list[float]` or `None` | Lambda values to try (incremental mode)         | `[0.001, 0.01, ..., 0.5]` |
+| `incremental_eps_y`     | `float`        | Max tolerated relative output-error increase               | `0.03`   |
+| `incremental_eps_w`     | `float`        | Min required relative weight-error decrease                | `0.10`   |
+| `incremental_initial_skip_ew_threshold` | `float` or `None` | Skip initial `lambda=0.0` candidate when Ew is too large | `0.3` |
+| `actorder`              | `bool`         | Reorder columns by activation magnitude                    | `False`  |
+| `device`                | `torch.device` | Device for computation                                     | `None`   |
+| `enable_clip_optimize`  | `bool`         | Enable Clip-Optimize initialization                        | `True`   |
+| `enable_clip_optimize_ep`| `bool`        | Enable Clip-Optimize with Error Propagation initialization | `False`  |
+| `enable_gptq`           | `bool`         | Enable GPTQ initialization                                 | `True`   |
+| `gptq`                  | `GPTQ` or `None` | Custom GPTQ instance for initial solution generation    | `None`   |
 
 ## Usage
 
@@ -73,6 +133,26 @@ runner = Runner(model_config=model_config, quantizer=jointq, qep=False)
 runner.run()
 ```
 
+### With incremental lambda
+
+```python
+jointq = JointQ(
+    bits=4,
+    group_size=128,
+    lambda_mode="incremental_lambda",
+)
+```
+
+### With identity regularization
+
+```python
+jointq = JointQ(
+    bits=4,
+    group_size=128,
+    regularization_mode="identity",
+)
+```
+
 ### With activation ordering
 
 ```python
@@ -85,10 +165,36 @@ jointq = JointQ(bits=4, group_size=128, actorder=True)
 jointq = JointQ(bits=4, symmetric=True, group_size=128)
 ```
 
+### With all initialization strategies
+
+```python
+jointq = JointQ(
+    bits=4,
+    group_size=128,
+    enable_clip_optimize=True,
+    enable_clip_optimize_ep=True,
+    enable_gptq=True,
+)
+```
+
+### Custom GPTQ parameters
+
+```python
+from onecomp.quantizer.gptq import GPTQ
+
+jointq = JointQ(
+    bits=4,
+    group_size=128,
+    gptq=GPTQ(wbits=4, groupsize=128, sym=False, mse=True, percdamp=0.05),
+)
+```
+
 ## Notes
 
 - JointQ requires GPU for computation (CUDA-based local search).
 - Group-wise quantization (`group_size > 0`) is recommended for accuracy.
-- The `batch_size` parameter controls memory usage: smaller values reduce peak GPU memory
-  at the cost of slower processing.
+  Set `group_size=None` for per-channel quantization.
 - JointQ currently supports dequantized-model evaluation only (not packed quantized inference).
+- Incremental lambda mode runs `quantize()` multiple times per layer (once per
+  lambda value until rejection), so quantization time increases compared to
+  fixed lambda mode.

onecomp/analyzer/__init__.py

@@ -6,6 +6,7 @@
 
 """
 
+from .quantization_error import plot_quantization_errors
 from .weight_outlier import (
     LayerOutlierStats,
     WeightOutlierAnalysis,
@@ -19,5 +20,6 @@
     "WeightOutlierAnalysis",
     "WeightOutlierAnalyzer",
     "analyze_weight_outliers",
+    "plot_quantization_errors",
     "save_weight_distribution_plots",
 ]