DOCSEARCHbyAlgolia
Submit Docs

SpinQuant Documentation

repository·main·Indexed 19 days ago

https://github.com/facebookresearch/spinquant

A quantization method for Large Language Models (LLMs) that utilizes learned rotations via Cayley optimization to remove outliers in activations. SpinQuant reduces the accuracy gap in low-bit quantization (e.g., W4A4KV4) and supports LLaMA-2 and LLaMA-3 models. It includes tools for optimizing rotation matrices, performing Post-Training Quantization (PTQ) using GPTQ or RTN, and exporting quantized models to ExecuTorch.

Tokens
2K
Snippets
7
Records
8
Agent score
65%

What's inside SpinQuant

  1. Install SpinQuant and its dependencies

    main

    To use SpinQuant, you need Python 3.9 and PyTorch >= 2.0 with CUDA support. You must also manually install the fast-hadamard-transform package from its source. Follow these steps:

    1. Install PyTorch with CUDA support.
    2. Install project requirements: pip install -r requirement.txt.
    3. Clone and install fast-hadamard-transform:
      git clone https://github.com/Dao-AILab/fast-hadamard-transform.git
cd fast-hadamard-transform
pip install .
    pip install -r requirement.txt
git clone https://github.com/Dao-AILab/fast-hadamard-transform.git
cd fast-hadamard-transform
pip install .

  2. Export quantized models to ExecuTorch

    main

    SpinQuant supports exporting quantized models to ExecuTorch to utilize specialized quantization kernels for real-time speedup. Currently, it supports 4-bit weight quantization (use group-size 256 for 8B models and 32 for smaller models) and 8-bit dynamic activation quantization.

    Use the following scripts for export:

    • scripts/31_optimize_rotation_executorch.sh $model_name
    • scripts/32_eval_ptq_executorch.sh $model_name
    bash scripts/31_optimize_rotation_executorch.sh $model_name
bash scripts/32_eval_ptq_executorch.sh $model_name

  3. Run PTQ evaluation with optimized rotation

    main

    After optimizing the rotation matrix, place the resulting matrix into your optimized_rotation_path and run the Post-Training Quantization (PTQ) evaluation script.

    Important Note on GPTQ: If you use the GPTQ method to quantize both weights and activations, you should optimize the rotation matrices with respect to a network where only activations are quantized. For example, optimize for W16A4KV4, then evaluate with W4A4KV4 using the optimized path from the W16 step.

    # Step 2: Run PTQ evaluation
bash scripts/2_eval_ptq.sh $model_name $w_bit $a_bit $kv_bit

  4. Optimize rotation matrices for LLaMA models

    main

    SpinQuant uses learned rotations to remove outliers and assist quantization. You can optimize rotation matrices using specific scripts depending on the model size.

    Note: You must set output_rotation_path, output_dir, logging_dir, and optimized_rotation_path to your own local paths. For gated models like meta-llama, provide your HuggingFace token via access_token.

    For LLaMA-2 7B/13B and LLaMA-3 8B:

    Use scripts/10_optimize_rotation.sh with arguments: $model_name $w_bit $a_bit $kv_bit.

    For LLaMA-2 70B and LLaMA-3 70B:

    Use scripts/11_optimize_rotation_fsdp.sh with arguments: $model_name $w_bit $a_bit $kv_bit.

    # Example for 4-bit weight, 4-bit activation, and 4-bit kv-cache on Llama-2-7b
bash scripts/10_optimize_rotation.sh meta-llama/Llama-2-7b 4 4 4

# Example for 70B models
bash scripts/11_optimize_rotation_fsdp.sh meta-llama/Llama-2-70b 4 4 4

  5. Export quantized models to ExecuTorch

    main

    When preparing a model for ExecuTorch, use the export_to_et flag. This triggers specific quantization and serialization logic:

    1. Specialized Layer Quantization: The lm_head and embed_tokens layers are quantized using 8-bit per-channel RTN quantization via gptq_utils.rtn_fwrd.
    2. Serialization: The model state dict is processed through write_model_llama (with num_shards=1) and then sanitized via sanitize_checkpoint_from_spinquant using the specified args.w_groupsize.

    This ensures the checkpoint format is compatible with the ExecuTorch runtime requirements.

  6. Reference: SpinQuant script arguments

    main

    The following arguments are used in the SpinQuant optimization and evaluation scripts:

    --input_model: The model name (or path to the weights)
--output_rotation_path: The local path we want to store the oprimized rotation matrix
--per_device_train_batch_size: The batch size for rotation optimization
--per_device_eval_batch_size: The batch size for PPL evaluation
--a_bits: The number of bits for activation quantization
--w_bits: The number of bits for weight quantization
--v_bits: The number of bits for value quantization
--k_bits: The number of bits for key quantization
--w_clip: Whether using the grid search to find best weight clipping range
--w_rtn: Whether we want to use round-to-nearest quantization. If not having `--w_rtn`, we are using GPTQ quantization.
--w_groupsize: The group size for group-wise weight quantization.
--rotate: Whether to rotate the model
--optimized_rotation_path: The checkpoint path of optimized rotation; Use random rotation if path is not given

  7. Prepare a model for rotation and quantization with `prepare_model`

    main

    The prepare_model function is the primary entrypoint for transforming a standard transformer model into a SpinQuant-ready model. It performs three main stages:

    1. Weight Rotation: Fuses layer norms and applies R3/R4 rotation to the model weights using apply_r3_r4.rotate_model.
    2. Activation Quantization Setup: Adds activation wrappers via quant_utils.add_actquant and configures Hadamard matrices for down_proj layers.
    3. Input Quantization: Configures bit-width, groupsize, symmetry, and clipping ratios for various projections (v_proj, o_proj, down_proj, etc.) and handles KV cache quantization (k_bits) by injecting wrappers after the RoPE function call.

    Note: If args.k_bits < 16, pre-RoPE quantization is currently not supported and will raise a NotImplementedError.

    # Usage pattern for prepare_model
# args must contain configuration for bits, groupsize, symmetry, etc.
model = prepare_model(args, model)

  8. Perform Post-Training Quantization (PTQ) with `ptq_model`

    main

    The ptq_model function is the primary entrypoint for applying SpinQuant's quantization pipeline to a model. It handles weight rotation, weight quantization (via GPTQ or RTN), and activation/KV cache quantization.

    Workflow Summary:

    1. Rotation: If args.rotate is enabled, it fuses layer norms and rotates the model weights using rotation_utils.rotate_model.
    2. Weight Quantization:
      • If args.w_bits < 16, it performs quantization.
      • If args.load_qmodel_path is provided, it loads an existing quantized/rotated model.
      • If args.w_rtn is false, it uses GPTQ quantization (gptq_utils.gptq_fwrd).
      • If args.w_rtn is true, it uses RTN quantization (gptq_utils.rtn_fwrd).
    3. Saving: If args.save_qmodel_path is set, it saves the state dict. If args.export_to_et is enabled, it uses write_model_llama and sanitize_checkpoint_from_spinquant to prepare the model for ExecuTorch.
    4. Activation & KV Cache Quantization: Configures quantization for activations (a_bits, v_bits) and KV cache (k_bits) by injecting wrappers into the model layers.
    def ptq_model(args, model, model_args=None):
    # ... implementation details ...
    return model