Merge branch 'export/issue-7-rotation-0-3-5' into 'export/v0-5-0'

Add rotation

See merge request onecomp/onecomp-lab!36

Author: FKKimura

CHANGELOG.md

@@ -15,6 +15,24 @@
   - Provides learning-based post-quantization fine-tuning for GPTQ-quantized models
   - Public API is exposed as `PostProcessLoraSFT`
 
+### New Feature: Rotation Preprocessing Pipeline (`onecomp/pre_process/`)
+
+SpinQuant/OstQuant-based rotation preprocessing that reduces quantization error by learning optimal rotation matrices before quantization. Supports Llama and Qwen3 architectures.
+
+- Added `prepare_rotated_model()` (`onecomp/pre_process/prepare_rotated_model.py`): End-to-end pipeline — model loading → rotation/scaling training → rotation application → saving
+  - Memory-optimized: moves model between CPU/GPU to reduce peak memory (e.g. Qwen3-32B: ~128GB → ~64GB)
+- Added `RotatedModelConfig` (`onecomp/rotated_model_config.py`): `ModelConfig` subclass that automatically registers Hadamard hooks on `down_proj` layers during `load_model()`
+- Added `onecomp/pre_process/` package:
+  - `train_rotation.py`: Training pipeline with `PreprocessManager` (R1/R2/S_* tensor management), HF `Trainer` subclass, `apply_preprocess_train` / `apply_preprocess_eval`
+  - `optimizer.py`: `SGDG` — SGD on the Stiefel manifold with Cayley-retraction orthogonal updates (ported from SpinQuant)
+  - `quant_models.py`: `WeightQuantizer` (RTN proxy) with per-channel / per-tensor / group-wise quantization; quantized decoder layers for Llama and Qwen3
+  - `rotation_utils.py`: `fuse_layer_norms`, `rotate_model`, `register_online_hadamard_hooks`
+  - `hadamard_utils.py`: Hadamard transform utilities and pre-computed matrices (ported from QuIP#)
+  - `modeling_llama.py` / `modeling_qwen3.py`: Custom `ForCausalLM` classes that propagate R1 through the forward pass during training
+  - `preprocess_args.py`: `TrainingArguments` subclass with SGDG-specific LR/momentum fields
+- Fixed `_PreprocessTrainer` to override `create_optimizer()` instead of `create_optimizer_and_scheduler()` for transformers >= 5.x compatibility (SGDG optimizer was silently replaced by AdamW)
+- Updated `Runner.save_dequantized_model()` and `Runner.save_quantized_model()` to warn when saving models loaded with additional preprocessing (e.g., Hadamard hooks)
+
 ### Added JointQ Quantizer
 
 - **Added new `JointQ` quantizer (`onecomp/quantizer/jointq/`)**
@@ -55,6 +73,8 @@
 - Added `example/post_process/example_lora_sft_knowledge.py`: Knowledge injection demo — teaches the quantized model about "OneCompression" via LoRA SFT and compares generation before/after
 - Added `example/post_process/onecomp_knowledge.jsonl`: Training data describing OneCompression for the knowledge injection example
 - Added `example/example_jointq.py`: JointQ 4-bit (groupsize=128) quantization example with dequantized model PPL evaluation
+- Added `example/pre_process/example_llama_preprocess_rtn.py`: Rotation preprocessing + RTN quantization (TinyLlama-1.1B)
+- Added `example/pre_process/example_preprocess_save_load.py`: Rotation preprocessing + GPTQ quantization → save → load → PPL verification
 
 ### Documentation
 
@@ -64,6 +84,11 @@
 - Updated `docs/api/runner.md` to include `create_quantized_model` and `save_quantized_model_pt`
 - Updated `docs/api/quantized_model_loader.md` to include `load_quantized_model_pt`
 - Updated `mkdocs.yml` navigation with new post-process pages
+- Added `docs/user-guide/pre-process.md`: Rotation preprocessing user guide covering workflow, key parameters, save/load, and limitations
+- Added `docs/api/pre_process.md`: API reference for `prepare_rotated_model` and `RotatedModelConfig`
+- Updated `docs/user-guide/examples.md` with rotation preprocessing code examples (RTN, GPTQ with save/load)
+- Updated `docs/api/index.md` with `RotatedModelConfig`, `prepare_rotated_model`, and `pre_process/` module structure
+- Updated `docs/index.md` Key Features with rotation preprocessing
 
 ### Tests
 

README.md

@@ -11,6 +11,7 @@ This package is currently under development (version 0) and may behave unstably.
 - **AutoBit**: Mixed-precision quantization with ILP-based bitwidth assignment. Automatically estimates the target bitwidth from available VRAM and assigns per-layer bitwidths to minimize quantization error under the memory budget.
 - **JointQ**: Joint quantization method that optimizes weight assignments and scale parameters simultaneously for improved quantization accuracy. Supports group-wise quantization (e.g., 4-bit, groupsize=128).
 - **LoRA SFT Post-Process**: Fine-tune quantized models with LoRA adapters for accuracy recovery or domain-specific knowledge injection. Supports SFT loss, teacher distillation, and intermediate block alignment.
+- **Rotation Preprocessing**: SpinQuant/OstQuant-based rotation preprocessing that reduces quantization error by learning optimal rotation matrices before quantization. Rotation/scaling matrices are absorbed into model weights, with online Hadamard hooks automatically registered at load time. Supports Llama and Qwen3 architectures.
 - (TBD)
 
 ## 🔧 Installation
@@ -142,11 +143,21 @@ uv run mkdocs serve
 
 Then open [http://127.0.0.1:8000](http://127.0.0.1:8000) in your browser.
 
-## 🚀 Example
-
-See [example/example1.py](./example/example1.py) and [example/example2.py](./example/example2.py) for step-by-step GPTQ + QEP examples.
-
-For AutoBit mixed-precision quantization, see [example/example3.py](./example/example3.py) and [example/example_auto_run.py](./example/example_auto_run.py).
+## 🚀 Examples
+
+| Category | Script | Description |
+|----------|--------|-------------|
+| Quantization | [example_gptq.py](./example/example_gptq.py) | GPTQ quantization |
+| | [example_qep_gptq.py](./example/example_qep_gptq.py) | GPTQ + QEP (error propagation) |
+| | [example_jointq.py](./example/example_jointq.py) | JointQ quantization |
+| | [example_autobit.py](./example/example_autobit.py) | AutoBit mixed-precision quantization |
+| | [example_auto_run.py](./example/example_auto_run.py) | AutoBit with automatic VRAM estimation |
+| Save / Load | [example_save_load.py](./example/example_save_load.py) | Save and load quantized models |
+| Rotation Preprocessing | [example_llama_preprocess_rtn.py](./example/pre_process/example_llama_preprocess_rtn.py) | Rotation preprocessing + RTN (TinyLlama) |
+| | [example_preprocess_save_load.py](./example/pre_process/example_preprocess_save_load.py) | Save and load rotation-preprocessed quantized models |
+| Post-Process | [example_lora_sft.py](./example/post_process/example_lora_sft.py) | LoRA SFT post-quantization fine-tuning |
+| | [example_lora_sft_knowledge.py](./example/post_process/example_lora_sft_knowledge.py) | LoRA SFT knowledge injection |
+| vLLM | [example_vllm_inference.py](./example/example_vllm_inference.py) | Serve quantized models with vLLM |
 
 ## 🔌 vLLM Inference
 

docs/api/index.md

@@ -21,8 +21,11 @@ The following are available directly from `import onecomp`:
 | `CQ`                   | Class    | CQ quantizer                             |
 | `QBB`                  | Class    | QBB quantizer                            |
 | `Onebit`               | Class    | 1-bit quantizer                          |
+| `RotatedModelConfig`   | Class    | ModelConfig for rotation-preprocessed models |
 | `QuantizedModelLoader` | Class    | Loader for saved quantized models        |
 | `load_quantized_model` | Function | Shortcut for `QuantizedModelLoader.load_quantized_model` |
+| `load_quantized_model_pt` | Function | Shortcut for `QuantizedModelLoader.load_quantized_model_pt` |
+| `prepare_rotated_model`| Function | Rotation preprocessing pipeline          |
 | `setup_logger`         | Function | Configure logging output                 |
 
 ## Module Structure
@@ -33,6 +36,9 @@ onecomp/
     cli.py                 # CLI entry point (onecomp command)
     __main__.py            # python -m onecomp support
     model_config.py        # ModelConfig class
+    rotated_model_config.py # RotatedModelConfig class
+    pre_process/           # Rotation preprocessing pipeline
+        prepare_rotated_model.py  # prepare_rotated_model()
     qep/                   # QEP module
         _qep_config.py     # QEPConfig dataclass
     quantizer/             # Quantizer implementations

docs/api/pre_process.md

@@ -0,0 +1,53 @@
+# Pre-Process (Rotation Preprocessing)
+
+Rotation preprocessing reduces quantization error by learning optimal rotation matrices
+(SpinQuant/OstQuant) and absorbing them into model weights before quantization.
+
+## `prepare_rotated_model`
+
+::: onecomp.pre_process.prepare_rotated_model.prepare_rotated_model
+    options:
+      show_source: false
+
+## `RotatedModelConfig`
+
+`ModelConfig` subclass for loading rotation-preprocessed models.
+Automatically registers Hadamard `forward_pre_hook` on `down_proj` layers.
+
+::: onecomp.rotated_model_config.RotatedModelConfig
+    options:
+      show_source: false
+
+## Workflow
+
+```
+┌─────────────────────────────────────────────────────────────┐
+│  Step 1: Rotation Preprocessing                             │
+│                                                             │
+│  ModelConfig ──► prepare_rotated_model() ──► RotatedModelConfig
+│                  (train rotation matrices,                  │
+│                   absorb into weights,                      │
+│                   save rotated model)                       │
+└──────────────────────────┬──────────────────────────────────┘
+                           │
+┌──────────────────────────▼──────────────────────────────────┐
+│  Step 2: Quantization                                       │
+│                                                             │
+│  RotatedModelConfig ──► Runner(quantizer=GPTQ/RTN/...) ──► run()
+│  (auto-registers            ──► save_quantized_model()      │
+│   Hadamard hooks)                                           │
+└──────────────────────────┬──────────────────────────────────┘
+                           │
+┌──────────────────────────▼──────────────────────────────────┐
+│  Step 3: Load                                               │
+│                                                             │
+│  load_quantized_model()                                     │
+│  (auto-detects "rotated: true" in config.json,              │
+│   registers Hadamard hooks automatically)                   │
+└─────────────────────────────────────────────────────────────┘
+```
+
+!!! note
+    The `wbits`, `groupsize`, and `sym` parameters passed to `prepare_rotated_model()`
+    control the RTN proxy used during rotation training. These values **must match**
+    the quantizer parameters used in Step 2.

docs/index.md

@@ -10,11 +10,12 @@ It implements state-of-the-art quantization algorithms including GPTQ, DBF, RTN,
 
 ## Key Features
 
-- **Multiple quantization algorithms** -- GPTQ, DBF (Double Binary Factorization), RTN (Round-To-Nearest), JointQ, QuIP, and more
-- **Quantization Error Propagation (QEP)** -- A novel method that adjusts weights before quantization to compensate for error propagation across layers
-- **Simple, unified API** -- Configure model, quantizer, and runner in a few lines of code
-- **Save/Load pipeline** -- Save quantized models in a format compatible with Hugging Face Transformers and vLLM
-- **Evaluation tools** -- Built-in perplexity and zero-shot accuracy benchmarks
+- **Quantization Error Propagation (QEP)** -- A post-training quantization method that corrects quantization errors by propagating them to subsequent layers, improving the accuracy of quantized LLMs. See [Arai & Ichikawa, NeurIPS 2025](https://openreview.net/forum?id=a3l3K9khbL) for details.
+- **vLLM Plugin Integration** -- Serve OneComp-quantized models with [vLLM](https://docs.vllm.ai/) via built-in plugins for DBF and Mixed-GPTQ quantization methods.
+- **AutoBit** -- Mixed-precision quantization with ILP-based bitwidth assignment. Automatically estimates the target bitwidth from available VRAM and assigns per-layer bitwidths to minimize quantization error under the memory budget.
+- **JointQ** -- Joint quantization method that optimizes weight assignments and scale parameters simultaneously for improved quantization accuracy. Supports group-wise quantization (e.g., 4-bit, groupsize=128).
+- **LoRA SFT Post-Process** -- Fine-tune quantized models with LoRA adapters for accuracy recovery or domain-specific knowledge injection. Supports SFT loss, teacher distillation, and intermediate block alignment.
+- **Rotation Preprocessing** -- SpinQuant/OstQuant-based rotation preprocessing that reduces quantization error by learning optimal rotation matrices before quantization. Rotation/scaling matrices are absorbed into model weights, with online Hadamard hooks automatically registered at load time. Supports Llama and Qwen3 architectures.
 
 ## Quick Example
 

docs/user-guide/examples.md

@@ -193,6 +193,81 @@ print(ppl_dict)
 # {'original': 5.47, 'GPTQ': 5.72, 'JointQ': 5.68}
 ```
 
+## Rotation Preprocessing + RTN
+
+Apply SpinQuant-style rotation preprocessing before quantization to reduce quantization error:
+
+```python
+from onecomp import ModelConfig, Runner, RTN, prepare_rotated_model, setup_logger
+
+setup_logger()
+
+model_config = ModelConfig(
+    model_id="TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T",
+    device="cuda:0",
+)
+
+# Step 1: Rotation preprocessing
+rotated_config = prepare_rotated_model(
+    model_config=model_config,
+    save_directory="./rotated_model",
+    seed=0,
+    wbits=3,
+    groupsize=-1,
+    sym=False,
+)
+
+# Step 2: Quantize the rotated model (wbits/groupsize/sym must match Step 1)
+rtn = RTN(wbits=3, groupsize=-1, sym=False)
+runner = Runner(model_config=rotated_config, quantizer=rtn)
+runner.run()
+
+original_ppl, dequantized_ppl, _ = runner.calculate_perplexity(
+    original_model=True, dequantized_model=True, quantized_model=False
+)
+print(f"Original model perplexity: {original_ppl}")
+print(f"Dequantized model perplexity: {dequantized_ppl}")
+```
+
+## Rotation Preprocessing + GPTQ with Save/Load
+
+Full pipeline including save and load of rotation-preprocessed quantized models:
+
+```python
+from onecomp import (
+    ModelConfig, Runner, GPTQ,
+    prepare_rotated_model, load_quantized_model, setup_logger,
+)
+
+setup_logger()
+
+# Step 1: Rotation preprocessing
+model_config = ModelConfig(
+    model_id="TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T",
+    device="cuda:0",
+)
+rotated_config = prepare_rotated_model(
+    model_config=model_config,
+    save_directory="./rotated_model",
+    seed=0,
+    wbits=4,
+    groupsize=128,
+)
+
+# Step 2: Quantize and save
+gptq = GPTQ(wbits=4, groupsize=128)
+runner = Runner(model_config=rotated_config, quantizer=gptq)
+runner.run()
+runner.save_quantized_model("./quantized_model")
+
+# Step 3: Load (Hadamard hooks are auto-registered via "rotated: true" in config.json)
+model, tokenizer = load_quantized_model("./quantized_model")
+```
+
+See [Pre-Process API](../api/pre_process.md) for full parameter documentation.
+
+---
+
 ## Saving and Loading Quantized Models
 
 ### Save the quantized model

docs/user-guide/pre-process.md

@@ -0,0 +1,92 @@
+# Pre-Process (Rotation Preprocessing)
+
+Rotation preprocessing applies SpinQuant/OstQuant-style rotation matrices to model weights
+before quantization, reducing quantization error. This is particularly effective for
+low-bit quantization (e.g., 3-bit).
+
+## Overview
+
+The rotation preprocessing pipeline:
+
+1. **Trains** rotation/scaling matrices using calibration data with an RTN quantization proxy
+2. **Absorbs** the learned matrices into model weights (fuses LayerNorms, rotates projections)
+3. **Registers** online Hadamard hooks on `down_proj` layers for inference correctness
+4. **Saves** the rotated model as a standard Hugging Face model directory
+
+The saved model can then be quantized with any quantizer (GPTQ, RTN, etc.) using the
+standard `Runner` pipeline.
+
+## Quick Start
+
+```python
+from onecomp import ModelConfig, Runner, GPTQ, prepare_rotated_model, setup_logger
+
+setup_logger()
+
+# Step 1: Rotation preprocessing
+model_config = ModelConfig(model_id="meta-llama/Llama-2-7b-hf", device="cuda:0")
+
+rotated_config = prepare_rotated_model(
+    model_config=model_config,
+    save_directory="./rotated_model",
+    wbits=4,
+    groupsize=128,
+)
+
+# Step 2: Quantize (wbits/groupsize/sym must match Step 1)
+gptq = GPTQ(wbits=4, groupsize=128)
+runner = Runner(model_config=rotated_config, quantizer=gptq)
+runner.run()
+```
+
+## Supported Architectures
+
+| Architecture | Status |
+|-------------|--------|
+| Llama       | Supported |
+| Qwen3       | Supported |
+
+## Key Parameters
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `rotation` | Apply rotation matrices (R1, R2) | `True` |
+| `scaling` | Apply scaling diagonals (S_*) | `False` |
+| `enable_training` | Train rotation matrices (vs. random init) | `True` |
+| `wbits` | RTN proxy bit-width (must match quantizer) | `4` |
+| `groupsize` | RTN proxy group size (must match quantizer) | `-1` |
+| `sym` | RTN proxy symmetric quantization | `False` |
+| `fp32_had` | Use FP32 for online Hadamard transform | `False` |
+| `seed` | Seed for rotation init and calibration data | `0` |
+
+!!! warning "Parameter matching"
+    The `wbits`, `groupsize`, and `sym` parameters **must match** the quantizer
+    used in Step 2. Mismatched values will degrade quantization quality because
+    the rotation matrices were optimized for different quantization settings.
+
+## Save and Load
+
+Rotation-preprocessed quantized models support the standard save/load API:
+
+```python
+# Save
+runner.save_quantized_model("./quantized_model")
+
+# Load (Hadamard hooks are auto-registered)
+from onecomp import load_quantized_model
+model, tokenizer = load_quantized_model("./quantized_model")
+```
+
+The saved `config.json` includes `"rotated": true` and `"fp32_had": false`,
+which `load_quantized_model()` uses to automatically register the required
+Hadamard hooks on `down_proj` layers.
+
+## Limitations
+
+- **vLLM inference is not supported.** vLLM kernels do not apply the online Hadamard
+  transform required by rotation-preprocessed models.
+- Only Llama and Qwen3 architectures are currently supported.
+
+## API Reference
+
+See [Pre-Process API](../api/pre_process.md) for full parameter documentation.

docs/user-guide/vllm-inference.md

@@ -10,6 +10,9 @@ The plugins are automatically registered via Python entry points when `onecomp`
 | DBF | `dbf` | 1-bit Double Binary Factorization. Uses GemLite kernels by default; set `ONECOMP_DBF_NAIVE_LINEAR=1` to use the naive fallback. |
 | Mixed-GPTQ | `mixed_gptq` | Per-layer mixed-bitwidth GPTQ. Automatically dispatches to Marlin or Exllama kernels based on bit-width and symmetry. |
 
+!!! warning "Rotation-preprocessed models are not supported"
+    Models quantized after rotation preprocessing (`prepare_rotated_model`) cannot be served with vLLM. vLLM kernels do not apply the online Hadamard transform on `down_proj` inputs that rotation-preprocessed models require for correct inference.
+
 ## Installation
 
 vLLM is available as an optional dependency: