orbitquant 0.1.6__tar.gz → 0.2.0__tar.gz

{orbitquant-0.1.6 → orbitquant-0.2.0}/.github/workflows/publish-pypi.yml

@@ -6,7 +6,7 @@ on:
       version:
         description: Version to publish
         required: true
-        default: "0.1.6"
+        default: "0.2.0"
 
 jobs:
   build:

orbitquant-0.2.0/PKG-INFO

@@ -0,0 +1,340 @@
+Metadata-Version: 2.4
+Name: orbitquant
+Version: 0.2.0
+Summary: Calibration-free OrbitQuant for transformer linear projections
+Project-URL: Homepage, https://github.com/iamwavecut/OrbitQuant
+Project-URL: Repository, https://github.com/iamwavecut/OrbitQuant
+Project-URL: Issues, https://github.com/iamwavecut/OrbitQuant/issues
+Project-URL: Paper, https://arxiv.org/abs/2607.02461
+Author: WaveCut
+License-Expression: Apache-2.0
+License-File: LICENSE
+Keywords: diffusers,diffusion,orbitquant,quantization,transformers
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Requires-Dist: huggingface-hub>=0.33
+Requires-Dist: numpy>=2.0
+Requires-Dist: packaging>=24.0
+Requires-Dist: safetensors>=0.5
+Requires-Dist: torch>=2.5
+Requires-Dist: tqdm>=4.66
+Provides-Extra: dev
+Requires-Dist: pytest-xdist>=3.8; extra == 'dev'
+Requires-Dist: pytest>=8.4; extra == 'dev'
+Requires-Dist: ruff>=0.12; extra == 'dev'
+Provides-Extra: eval
+Requires-Dist: imageio-ffmpeg>=0.6; extra == 'eval'
+Requires-Dist: imageio>=2.37; extra == 'eval'
+Requires-Dist: pillow>=11.0; extra == 'eval'
+Provides-Extra: hf
+Requires-Dist: accelerate>=1.8; extra == 'hf'
+Requires-Dist: diffusers>=0.35; extra == 'hf'
+Requires-Dist: transformers>=4.53; extra == 'hf'
+Provides-Extra: kernels
+Requires-Dist: kernels>=0.16; extra == 'kernels'
+Requires-Dist: triton>=3.5; (platform_system == 'Linux') and extra == 'kernels'
+Description-Content-Type: text/markdown
+
+# OrbitQuant
+
+OrbitQuant is a calibration-free post-training quantizer for transformer linear
+projections. It implements the method from
+[OrbitQuant: Data-Agnostic Quantization for Image and Video Diffusion Transformers](https://arxiv.org/abs/2607.02461)
+and exposes it through Hugging Face Transformers, Diffusers, and a direct
+PyTorch API.
+
+The implementation is clean-room and Apache-2.0 licensed.
+
+## Features
+
+- Automatic coverage of registered linear-compatible modules in transformer
+  backbones, independent of model class or modality.
+- Built-in support for `torch.nn.Linear` and Hugging Face `Conv1D` projections.
+- Public adapters for custom `F.linear`-equivalent module types and transposed
+  weight layouts.
+- RPBH rotation, exact unit-sphere Lloyd-Max codebooks, packed 2/3/4/6/8-bit
+  weights, and online activation quantization without calibration data.
+- Model-specific policies for paper-sensitive AdaLN and output-layer handling.
+- Compact `safetensors` artifacts and Hugging Face `save_pretrained()` /
+  `from_pretrained()` integration.
+- Packed-weight CUDA, Triton, and Metal inference paths that avoid a full
+  dequantized weight matrix.
+
+Embeddings, timestep modules, task heads, and common final projections are
+kept in source precision by default. Every automatic decision is available as
+a machine-readable inventory before quantization.
+
+## Install
+
+```bash
+pip install "orbitquant[hf]"
+```
+
+Install the CUDA/Triton and local-kernel loader dependencies with:
+
+```bash
+pip install "orbitquant[hf,kernels]"
+```
+
+## Quantize A Transformers Model
+
+Importing `orbitquant` registers the backend with supported Transformers and
+Diffusers versions. The default `target_policy="auto"` selects a known
+paper policy where applicable and otherwise uses the universal policy.
+
+```python
+import torch
+import orbitquant
+from transformers import AutoModelForCausalLM
+
+config = orbitquant.recipe(
+    "w4a4",
+    runtime_mode="auto_fused",
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    "your-org/your-transformer",
+    dtype=torch.bfloat16,
+    quantization_config=config,
+)
+model.save_pretrained("./your-transformer-orbitquant-w4a4")
+```
+
+Load the packed model after importing the backend:
+
+```python
+import orbitquant
+from transformers import AutoModelForCausalLM
+
+model = AutoModelForCausalLM.from_pretrained(
+    "./your-transformer-orbitquant-w4a4",
+    device_map="auto",
+)
+```
+
+Named recipes are `w4a4`, `w3a3`, `w2a4`, `w2a3`, and `w4a6`. They create a
+normal `OrbitQuantConfig`, so every field can be overridden.
+
+## Inspect Coverage
+
+Inspect a model before replacing modules:
+
+```python
+from orbitquant import inspect_linear_module_policy, recipe
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("your-org/your-transformer")
+report = inspect_linear_module_policy(model, recipe("w4a4"))
+
+print(report["action_counts"])
+print(report["quantized_modules"])
+print(report["skipped_modules"])
+print(report["unsupported_linear_modules"])
+```
+
+The universal policy quantizes every registered linear-compatible module except
+known embeddings, timestep modules, task/output heads, and explicit skips. It
+does not depend on names such as `layers`, `blocks`, or a particular model
+class.
+
+Use `modules_to_convert` as an allowlist and define AdaLN/skips with exact names,
+substrings, or glob patterns:
+
+```python
+from orbitquant import OrbitQuantConfig
+
+config = OrbitQuantConfig(
+    modules_to_convert=["backbone.*.projection"],
+    modules_to_use_adaln=["backbone.*.modulation"],
+    modules_to_not_convert=["*.sensitive_output"],
+)
+```
+
+Explicit dtype overrides remain available through `modules_dtype_dict`.
+
+## Quantize An Instantiated Module
+
+For ordinary PyTorch models or frameworks that do not use Hugging Face loading
+hooks:
+
+```python
+from orbitquant import quantize_model, recipe
+
+summary = quantize_model(
+    model,
+    recipe("w4a4"),
+    quantization_device="cuda",
+)
+print(summary.quantized_modules)
+```
+
+The replacement supports arbitrary leading dimensions and treats the final
+dimension as `in_features`, including sequence, image-token, and video-token
+layouts.
+
+## Custom Linear Modules
+
+Register a module whose forward operation is equivalent to `F.linear`. The
+adapter describes only its source weight layout and feature attributes:
+
+```python
+from orbitquant import register_linear_adapter
+
+register_linear_adapter(
+    MyLinear,
+    weight_layout="in_out",
+    in_features_attr="input_size",
+    out_features_attr="output_size",
+)
+```
+
+OrbitQuant stores every replacement in canonical `[out_features, in_features]`
+layout. Modules with additional routing, tensor-parallel communication, sparse
+expert selection, or non-linear forward semantics need an architecture-aware
+adapter; the inspection report lists unregistered linear candidates instead of
+silently replacing them.
+
+## Diffusers
+
+Quantize the denoiser component of a pipeline:
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+from orbitquant import quantize_pipeline, recipe, save_quantized_pipeline_component
+
+pipe = DiffusionPipeline.from_pretrained(
+    "black-forest-labs/FLUX.2-klein-4B",
+    torch_dtype=torch.bfloat16,
+)
+config = recipe("w4a4", target_policy="flux2")
+summary = quantize_pipeline(
+    pipe,
+    config,
+    component="transformer",
+    quantization_device="cuda",
+)
+save_quantized_pipeline_component(
+    pipe,
+    "./flux2-klein-orbitquant-w4a4",
+    config=config,
+    component="transformer",
+    source_model_id="black-forest-labs/FLUX.2-klein-4B",
+    summary=summary,
+)
+```
+
+Published FLUX, Z-Image, and Wan repositories are compact Diffusers component
+artifacts. Their model cards contain the matching pipeline code and native
+generation settings.
+
+Load a published component artifact together with its recorded source pipeline:
+
+```python
+import torch
+from huggingface_hub import snapshot_download
+from orbitquant import load_quantized_pipeline_from_artifact
+
+artifact_dir = snapshot_download(
+    "WaveCut/FLUX.1-schnell-OrbitQuant-W4A4",
+    repo_type="model",
+)
+pipe = load_quantized_pipeline_from_artifact(
+    artifact_dir,
+    torch_dtype=torch.bfloat16,
+    runtime_mode="auto_fused",
+)
+```
+
+## Packed Runtime
+
+`runtime_mode="auto_fused"` is the default:
+
+| Device | Dispatch |
+| --- | --- |
+| CUDA | Native packed CUDA package, then Triton packed matmul |
+| MPS | Native packed Metal package |
+| CPU | PyTorch reference path |
+
+CUDA and MPS do not silently materialize a full BF16/FP16 weight matrix in
+`auto_fused`. If no packed backend is available, the error includes the missing
+backend and installation guidance.
+
+Use the explicit reference path for compatibility or numerical debugging:
+
+```python
+config = orbitquant.recipe("w4a4", runtime_mode="dequant_bf16")
+```
+
+Build the ABI3 native package locally without Kernel Hub:
+
+```bash
+cd native-kernels/orbitquant-packed-matmul
+nix --option sandbox relaxed run .#build-and-copy -L
+export PYTHONPATH="$PWD/build/<matching-torch-backend-platform-variant>:$PYTHONPATH"
+```
+
+The variant must match the Torch minor version, CUDA or Metal backend, C++ ABI,
+architecture, and operating system. See
+[`docs/kernel-audit.md`](docs/kernel-audit.md) for tested shapes, benchmark
+methodology, and local package verification.
+
+## Validated Architecture Coverage
+
+The integration suite instantiates and inventories encoder-only, decoder-only,
+encoder-decoder, causal LM, and vision transformer families:
+
+| Family | Projection type |
+| --- | --- |
+| BERT | `torch.nn.Linear` |
+| GPT-2 | Hugging Face `Conv1D` with transposed source weights |
+| Llama | `torch.nn.Linear`, including GQA projections |
+| T5 | encoder and decoder `torch.nn.Linear` projections |
+| ViT | vision transformer `torch.nn.Linear` projections |
+
+The paper-aligned release artifacts remain FLUX.1-schnell, Z-Image-Turbo, and
+Wan 2.1 T2V. FLUX.2 Klein is an additional validated diffusion target.
+
+Architecture coverage means the model can be discovered, quantized, executed,
+saved, and restored without model-name-specific code. It does not guarantee a
+quality-preserving bit setting. OrbitQuant was evaluated in the paper on image
+and video diffusion transformers; language and classification models can be
+more sensitive, and their quality must be measured before publishing a
+checkpoint. The library exposes module overrides for such recipes but does not
+silently substitute a different quantization algorithm.
+
+## Method Conformance
+
+The implementation follows the paper's shared data-agnostic basis:
+
+- RPBH permutation, Rademacher signs, block FWHT, and orthonormal scaling.
+- Offline folded weight rotation with BF16 row norms and quantized unit
+  directions.
+- Online per-token norm, normalized activation rotation, nearest-centroid
+  quantization, and rescaling.
+- One fixed codebook per `(input dimension, bit width, algorithm version)` and
+  no prompt, timestep, or calibration statistics.
+- INT4 group-64 RTN for model policies that identify dynamic AdaLN projections.
+
+The detailed requirement matrix is in
+[`docs/paper-methodology-audit.md`](docs/paper-methodology-audit.md).
+
+## Development
+
+```bash
+uv sync --extra hf --extra dev
+uv run pytest -q
+uv run ruff check .
+```
+
+## License
+
+OrbitQuant is licensed under Apache-2.0. Model artifacts retain the license and
+provenance of their source checkpoints.

orbitquant-0.2.0/README.md

@@ -0,0 +1,296 @@
+# OrbitQuant
+
+OrbitQuant is a calibration-free post-training quantizer for transformer linear
+projections. It implements the method from
+[OrbitQuant: Data-Agnostic Quantization for Image and Video Diffusion Transformers](https://arxiv.org/abs/2607.02461)
+and exposes it through Hugging Face Transformers, Diffusers, and a direct
+PyTorch API.
+
+The implementation is clean-room and Apache-2.0 licensed.
+
+## Features
+
+- Automatic coverage of registered linear-compatible modules in transformer
+  backbones, independent of model class or modality.
+- Built-in support for `torch.nn.Linear` and Hugging Face `Conv1D` projections.
+- Public adapters for custom `F.linear`-equivalent module types and transposed
+  weight layouts.
+- RPBH rotation, exact unit-sphere Lloyd-Max codebooks, packed 2/3/4/6/8-bit
+  weights, and online activation quantization without calibration data.
+- Model-specific policies for paper-sensitive AdaLN and output-layer handling.
+- Compact `safetensors` artifacts and Hugging Face `save_pretrained()` /
+  `from_pretrained()` integration.
+- Packed-weight CUDA, Triton, and Metal inference paths that avoid a full
+  dequantized weight matrix.
+
+Embeddings, timestep modules, task heads, and common final projections are
+kept in source precision by default. Every automatic decision is available as
+a machine-readable inventory before quantization.
+
+## Install
+
+```bash
+pip install "orbitquant[hf]"
+```
+
+Install the CUDA/Triton and local-kernel loader dependencies with:
+
+```bash
+pip install "orbitquant[hf,kernels]"
+```
+
+## Quantize A Transformers Model
+
+Importing `orbitquant` registers the backend with supported Transformers and
+Diffusers versions. The default `target_policy="auto"` selects a known
+paper policy where applicable and otherwise uses the universal policy.
+
+```python
+import torch
+import orbitquant
+from transformers import AutoModelForCausalLM
+
+config = orbitquant.recipe(
+    "w4a4",
+    runtime_mode="auto_fused",
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    "your-org/your-transformer",
+    dtype=torch.bfloat16,
+    quantization_config=config,
+)
+model.save_pretrained("./your-transformer-orbitquant-w4a4")
+```
+
+Load the packed model after importing the backend:
+
+```python
+import orbitquant
+from transformers import AutoModelForCausalLM
+
+model = AutoModelForCausalLM.from_pretrained(
+    "./your-transformer-orbitquant-w4a4",
+    device_map="auto",
+)
+```
+
+Named recipes are `w4a4`, `w3a3`, `w2a4`, `w2a3`, and `w4a6`. They create a
+normal `OrbitQuantConfig`, so every field can be overridden.
+
+## Inspect Coverage
+
+Inspect a model before replacing modules:
+
+```python
+from orbitquant import inspect_linear_module_policy, recipe
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("your-org/your-transformer")
+report = inspect_linear_module_policy(model, recipe("w4a4"))
+
+print(report["action_counts"])
+print(report["quantized_modules"])
+print(report["skipped_modules"])
+print(report["unsupported_linear_modules"])
+```
+
+The universal policy quantizes every registered linear-compatible module except
+known embeddings, timestep modules, task/output heads, and explicit skips. It
+does not depend on names such as `layers`, `blocks`, or a particular model
+class.
+
+Use `modules_to_convert` as an allowlist and define AdaLN/skips with exact names,
+substrings, or glob patterns:
+
+```python
+from orbitquant import OrbitQuantConfig
+
+config = OrbitQuantConfig(
+    modules_to_convert=["backbone.*.projection"],
+    modules_to_use_adaln=["backbone.*.modulation"],
+    modules_to_not_convert=["*.sensitive_output"],
+)
+```
+
+Explicit dtype overrides remain available through `modules_dtype_dict`.
+
+## Quantize An Instantiated Module
+
+For ordinary PyTorch models or frameworks that do not use Hugging Face loading
+hooks:
+
+```python
+from orbitquant import quantize_model, recipe
+
+summary = quantize_model(
+    model,
+    recipe("w4a4"),
+    quantization_device="cuda",
+)
+print(summary.quantized_modules)
+```
+
+The replacement supports arbitrary leading dimensions and treats the final
+dimension as `in_features`, including sequence, image-token, and video-token
+layouts.
+
+## Custom Linear Modules
+
+Register a module whose forward operation is equivalent to `F.linear`. The
+adapter describes only its source weight layout and feature attributes:
+
+```python
+from orbitquant import register_linear_adapter
+
+register_linear_adapter(
+    MyLinear,
+    weight_layout="in_out",
+    in_features_attr="input_size",
+    out_features_attr="output_size",
+)
+```
+
+OrbitQuant stores every replacement in canonical `[out_features, in_features]`
+layout. Modules with additional routing, tensor-parallel communication, sparse
+expert selection, or non-linear forward semantics need an architecture-aware
+adapter; the inspection report lists unregistered linear candidates instead of
+silently replacing them.
+
+## Diffusers
+
+Quantize the denoiser component of a pipeline:
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+from orbitquant import quantize_pipeline, recipe, save_quantized_pipeline_component
+
+pipe = DiffusionPipeline.from_pretrained(
+    "black-forest-labs/FLUX.2-klein-4B",
+    torch_dtype=torch.bfloat16,
+)
+config = recipe("w4a4", target_policy="flux2")
+summary = quantize_pipeline(
+    pipe,
+    config,
+    component="transformer",
+    quantization_device="cuda",
+)
+save_quantized_pipeline_component(
+    pipe,
+    "./flux2-klein-orbitquant-w4a4",
+    config=config,
+    component="transformer",
+    source_model_id="black-forest-labs/FLUX.2-klein-4B",
+    summary=summary,
+)
+```
+
+Published FLUX, Z-Image, and Wan repositories are compact Diffusers component
+artifacts. Their model cards contain the matching pipeline code and native
+generation settings.
+
+Load a published component artifact together with its recorded source pipeline:
+
+```python
+import torch
+from huggingface_hub import snapshot_download
+from orbitquant import load_quantized_pipeline_from_artifact
+
+artifact_dir = snapshot_download(
+    "WaveCut/FLUX.1-schnell-OrbitQuant-W4A4",
+    repo_type="model",
+)
+pipe = load_quantized_pipeline_from_artifact(
+    artifact_dir,
+    torch_dtype=torch.bfloat16,
+    runtime_mode="auto_fused",
+)
+```
+
+## Packed Runtime
+
+`runtime_mode="auto_fused"` is the default:
+
+| Device | Dispatch |
+| --- | --- |
+| CUDA | Native packed CUDA package, then Triton packed matmul |
+| MPS | Native packed Metal package |
+| CPU | PyTorch reference path |
+
+CUDA and MPS do not silently materialize a full BF16/FP16 weight matrix in
+`auto_fused`. If no packed backend is available, the error includes the missing
+backend and installation guidance.
+
+Use the explicit reference path for compatibility or numerical debugging:
+
+```python
+config = orbitquant.recipe("w4a4", runtime_mode="dequant_bf16")
+```
+
+Build the ABI3 native package locally without Kernel Hub:
+
+```bash
+cd native-kernels/orbitquant-packed-matmul
+nix --option sandbox relaxed run .#build-and-copy -L
+export PYTHONPATH="$PWD/build/<matching-torch-backend-platform-variant>:$PYTHONPATH"
+```
+
+The variant must match the Torch minor version, CUDA or Metal backend, C++ ABI,
+architecture, and operating system. See
+[`docs/kernel-audit.md`](docs/kernel-audit.md) for tested shapes, benchmark
+methodology, and local package verification.
+
+## Validated Architecture Coverage
+
+The integration suite instantiates and inventories encoder-only, decoder-only,
+encoder-decoder, causal LM, and vision transformer families:
+
+| Family | Projection type |
+| --- | --- |
+| BERT | `torch.nn.Linear` |
+| GPT-2 | Hugging Face `Conv1D` with transposed source weights |
+| Llama | `torch.nn.Linear`, including GQA projections |
+| T5 | encoder and decoder `torch.nn.Linear` projections |
+| ViT | vision transformer `torch.nn.Linear` projections |
+
+The paper-aligned release artifacts remain FLUX.1-schnell, Z-Image-Turbo, and
+Wan 2.1 T2V. FLUX.2 Klein is an additional validated diffusion target.
+
+Architecture coverage means the model can be discovered, quantized, executed,
+saved, and restored without model-name-specific code. It does not guarantee a
+quality-preserving bit setting. OrbitQuant was evaluated in the paper on image
+and video diffusion transformers; language and classification models can be
+more sensitive, and their quality must be measured before publishing a
+checkpoint. The library exposes module overrides for such recipes but does not
+silently substitute a different quantization algorithm.
+
+## Method Conformance
+
+The implementation follows the paper's shared data-agnostic basis:
+
+- RPBH permutation, Rademacher signs, block FWHT, and orthonormal scaling.
+- Offline folded weight rotation with BF16 row norms and quantized unit
+  directions.
+- Online per-token norm, normalized activation rotation, nearest-centroid
+  quantization, and rescaling.
+- One fixed codebook per `(input dimension, bit width, algorithm version)` and
+  no prompt, timestep, or calibration statistics.
+- INT4 group-64 RTN for model policies that identify dynamic AdaLN projections.
+
+The detailed requirement matrix is in
+[`docs/paper-methodology-audit.md`](docs/paper-methodology-audit.md).
+
+## Development
+
+```bash
+uv sync --extra hf --extra dev
+uv run pytest -q
+uv run ruff check .
+```
+
+## License
+
+OrbitQuant is licensed under Apache-2.0. Model artifacts retain the license and
+provenance of their source checkpoints.