| --- |
| language: en |
| library_name: agiws-neural-quant |
| license: apache-2.0 |
| tags: |
| - quantization |
| - int4 |
| - int8 |
| - fp8 |
| - fp4 |
| - nf4 |
| - bitnet |
| - qat |
| - pytorch |
| --- |
| |
| # AGIWSNeuralQuant |
|
|
| Universal neural network quantization library — 53 formats, weights + activations, QAT, residual codebooks, per-head scales. |
|
|
| ## What |
|
|
| Quantize ANY PyTorch model (vision encoders, LLMs, TTS, image/video generation) via ONE parameterized `Quantizer`: |
|
|
| - **INT** — int2/3/4/6/8, per-tensor/per-channel/per-group/per-head scales, symmetric/asymmetric |
| - **MX/NV formats** — MXFP4/MXFP6/MXFP8 (E8M0 block scale), NVFP4/NVFP6/NVFP8 (FP8 scale), MXINT2/4/6/8 |
| - **FP4/FP6/FP8** — E2M1, E3M2/E2M3, E4M3/E5M2 floating-point with LUT-based dequant |
| - **NF4** — QLoRA NormalFloat 4-bit + double quantization |
| - **GGUF k-quants** — q2_k/q3_k/q4_k/q5_k/q6_k/q8_0/q4_0 (super-block layout) |
| - **BitNet 1.58** — ternary {-1,0,+1}, binary {+1,-1} |
| - **Codebook/VQ** — kmeans/fixed codebook, vector quantization, AQLM, residual multi-level codebooks |
| - **PTQ** — GPTQ (Hessian compensation), AWQ (activation-aware), SmoothQuant, LLM.int8 outlier, QuIP (rotation) |
| - **QAT** — learnable latent weights + scale via STE, learnable codebooks, dual-path distillation |
| - **Pruning** — structural (REAP/REAM), unstructured, magnitude-based |
| - **Activations** — ALL formats quantize activations, per-token/per-group/per-channel scale modes |
| - **per-head** — attention-specific scale per head (Q/K/V/O projections) |
| |
| ## Install |
| |
| ```bash |
| pip install git+https://huggingface.co/agiws/NeuralQuant |
| ``` |
| |
| ## Usage |
| |
| ```python |
| from agiws_neural_quant import quantize_model, save_model, load_model |
|
|
| # Quantize any PyTorch model in-place (53 formats via presets) |
| quantize_model(model, format="int4") # 4-bit per-group |
| quantize_model(model, format="mxfp4") # MXFP4 (E8M0 block scale) |
| quantize_model(model, format="nvfp4") # NVFP4 (FP8 scale, weights+activations) |
| quantize_model(model, format="ternary") # BitNet 1.58 |
| quantize_model(model, format="codebook", |
| residual_levels=3, # residual codebooks cascade |
| learnable=True) # QAT via STE |
| |
| # per-head for attention projections |
| quantize_model(model, format="int4", |
| per_head_modules=["q_proj","k_proj","v_proj","o_proj"], |
| num_heads=32, head_dim=128) |
| |
| # Save / load quantized model (preserves all format metadata) |
| save_model(model, "quantized.pt") |
| model2 = load_model("quantized.pt") |
| |
| # QAT training |
| from agiws_neural_quant import UnifiedQATWrapper, strip_latent |
| qat = UnifiedQATWrapper(quantized_module) |
| loss = qat.distillation_loss(x) # MSE(student, teacher.detach()) |
| # ... optimizer.step() ... |
| strip_latent(quantized_module) # bake to frozen inference |
| ``` |
| |
| ## Structure |
| |
| ``` |
| src/agiws_neural_quant/ |
| base.py — QuantizedWeight, QuantizedActivation, QuantizedModule (chunked dequant, dual-path, QAT) |
| quantizer.py — unified Quantizer (53 formats as parameter configurations) |
| presets.py — FORMAT_PRESETS dict (53 presets) + get_preset(format, **overrides) |
| dispatch.py — quantize_model(), save_model(), load_model(), walker |
| ssm_patterns.py — SSM/KDA exclusion patterns (Mamba/RWKV/linear-attention aware) |
| training_unified.py — UnifiedQATWrapper, dual_path_loss, strip_latent |
| training/ste.py — STEQuantize, STECodebook (Straight-Through Estimator) |
| ternary.py — ternarize_tensor (BitNet 1.58 primitives) |
| nf4/ — NF4 LUT, quantize/dequantize, double quantization (QLoRA) |
| fp4/ — FP4 E2M1 LUT, pack/unpack, E8M0 LUT |
| fp6/ — FP6 E3M2/E2M3 LUT, quantize/dequantize, pack/unpack |
| fp8/ — FP8 E4M3/E5M2 LUT, quantize/dequantize |
| kquant/ — GGUF k-quants (super-block layout) |
| converters/ — universal file-to-file converter (stream, low memory) |
| universal.py — convert_model(input, output): safetensors/GGUF/.pt <-> .pt/safetensors/GGUF |
| safetensors_io.py — stream read/write safetensors (lazy, low memory) |
| gguf_reader.py — GGUF binary format parser + k-quant decode |
| hf_nvfp4.py — NVFP4 safetensors -> QuantizedWeight |
| analysis.py — layer analysis (per-layer quantization suitability) |
| cache.py — teacher-cache for distillation |
| extract.py — checkpoint extraction (vision encoders, submodules) |
| tests/ — 169 tests (roundtrip, chunked, dual-path, QAT, SSM, residual, per-head, validation, converter) |
| scripts/ — auxiliary scripts |
| docs/ |
| ``` |
| |
| ## Principles |
| |
| - ONE parameterized `Quantizer` — all formats are configurations, not separate classes |
| - Weights AND activations quantized for ALL formats (W8A8, W4A4, NVFP4) |
| - Chunked dequant for minimal VRAM (adaptive chunk_size at runtime) |
| - Dual-path: student + teacher QuantizedModule for cross-quantization distillation |
| - Save/load preserves all format metadata (buffers + meta + quantizer config + dual-path + teacher) |
| - Supports any nn.Module: Linear, Conv1d/2d/3d, ConvTranspose, Embedding, LayerNorm, Bilinear |
| - Parameter validation: invalid combinations raise clear errors (not silent fallback) |
| - SSM-aware: excludes critical projections (q_proj/k_proj/b_proj) from quantization for Mamba/KDA/RWKV |
| |
| ## Converter (file-to-file, low memory) |
| |
| ```python |
| from agiws_neural_quant import convert_model |
| |
| # NVFP4 safetensors (GLM-5.2) -> NeuralQuant .pt (preserves quantized layout) |
| convert_model("glm5.2-nvfp4.safetensors", "glm5.2.pt") |
| |
| # fp16 safetensors -> NeuralQuant .pt with int4 quantization |
| convert_model("model.safetensors", "model_int4.pt", quant_format="int4") |
| |
| # GGUF k-quant -> safetensors (dequantized to fp16) |
| convert_model("model.gguf", "model_fp16.safetensors", output_format="safetensors") |
| |
| # NeuralQuant .pt -> safetensors (dequantized) |
| convert_model("model.pt", "model_out.safetensors", output_format="safetensors") |
| ``` |
| |
| Processes one tensor at a time (peak memory ~2x largest tensor, not the full model). |
| Supports sharded safetensors (index.json + multiple shards). |
| |
| ## Status |
| |
| - 53 quantization presets, all validated (roundtrip tested) |
| - Universal file-to-file converter (safetensors/GGUF/.pt, stream, low memory) |
| - 169 tests PASS |
| - Version 0.2.0 |
