### Run SpinQuant Optimization CLI

Source: https://context7.com/facebookresearch/spinquant/llms.txt

This command initiates the rotation optimization and quantization process for a specified Llama model. It configures model parameters, quantization settings for weights and activations, and specifies export paths for the resulting model.

```bash
torchrun --nnodes=1 --nproc_per_node=8 optimize_rotation.py \
    --input_model "meta-llama/Llama-2-7b" \
    --access_token "hf_xxxxx" \
    --output_rotation_path "./rotations" \
    --optimized_rotation_path "./R.bin" \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 4 \
    --model_max_length 2048 \
    --learning_rate 1.5 \
    --max_steps 100 \
    --bf16 True \
    --w_bits 4 \
    --w_groupsize 128 \
    --w_asym \
    --w_clip \
    --w_rtn \
    --a_bits 4 \
    --a_groupsize -1 \
    --a_asym \
    --a_clip_ratio 1.0 \
    --k_bits 4 \
    --v_bits 4 \
    --k_groupsize 128 \
    --v_groupsize 128 \
    --k_asym \
    --v_asym \
    --rotate \
    --rotate_mode "hadamard" \
    --fp32_had \
    --nsamples 128 \
    --percdamp 0.01 \
    --act_order \
    --save_qmodel_path "./model.pth" \
    --export_to_et
```

--------------------------------

### Load WikiText-2 Calibration and Evaluation Data

Source: https://context7.com/facebookresearch/spinquant/llms.txt

Uses the get_wikitext2 utility to prepare data for GPTQ calibration or perplexity evaluation. It supports switching between training mode for calibration samples and evaluation mode for full test sets.

```python
from utils.data_utils import get_wikitext2
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b")

# Get calibration samples for GPTQ
trainloader = get_wikitext2(nsamples=128, seed=0, seqlen=2048, tokenizer=tokenizer, eval_mode=False)

# Get evaluation data for perplexity
testenc = get_wikitext2(seqlen=2048, tokenizer=tokenizer, eval_mode=True)
```

--------------------------------

### WikiText-2 Data Loading

Source: https://context7.com/facebookresearch/spinquant/llms.txt

This section details how to load and prepare the WikiText-2 dataset for either GPTQ calibration or perplexity evaluation using the `get_wikitext2` function.

```APIDOC
## WikiText-2 Data Loading

The `get_wikitext2` function prepares calibration data for GPTQ quantization or evaluation data for perplexity measurement.

### Usage for GPTQ Calibration

```python
from utils.data_utils import get_wikitext2
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b")

trainloader = get_wikitext2(
    nsamples=128,    # Number of calibration samples
    seed=0,          # Random seed for reproducibility
    seqlen=2048,     # Sequence length
    tokenizer=tokenizer,
    eval_mode=False  # Returns list of (input_ids, targets) tuples
)

# trainloader: [(input_ids, targets), ...] with 128 samples
for inp, tar in trainloader[:3]:
    print(f"Input shape: {inp.shape}, Target shape: {tar.shape}")
# Expected output:
# Input shape: torch.Size([1, 2048]), Target shape: torch.Size([1, 2048])
```

### Usage for Perplexity Evaluation

```python
from utils.data_utils import get_wikitext2
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b")

testenc = get_wikitext2(
    seqlen=2048,
    tokenizer=tokenizer,
    eval_mode=True   # Returns full tokenized test set
)
# testenc.input_ids shape: torch.Size([1, ~250000])
```

### Parameters for `get_wikitext2`

- **nsamples** (int) - Optional - Number of calibration samples to generate.
- **seed** (int) - Optional - Random seed for reproducibility.
- **seqlen** (int) - Required - Sequence length for the data.
- **tokenizer** (object) - Required - Hugging Face tokenizer instance.
- **eval_mode** (bool) - Required - If True, returns the full tokenized test set; otherwise, returns calibration samples.
```

--------------------------------

### QKRotationWrapper for K-Cache Quantization

Source: https://context7.com/facebookresearch/spinquant/llms.txt

This section explains the `QKRotationWrapper`, which enhances K-cache quantization by applying Hadamard transforms and quantization to keys after RoPE.

```APIDOC
## QKRotationWrapper for K-Cache Quantization

The `QKRotationWrapper` applies Hadamard transform and quantization to keys after RoPE (Rotary Position Embedding) for efficient KV-cache quantization.

### Usage

```python
from eval_utils.rotation_utils import add_qk_rotation_wrapper_after_function_call_in_forward

# Assume model and model.config are initialized
# model = ... (your initialized model)

# Configure K-cache quantization parameters
k_quant_config = {
    "k_bits": 4,           # Bits for key cache
    "k_groupsize": 128,    # Group size (128 or -1 for token-wise)
    "k_sym": False,        # Asymmetric quantization
    "k_clip_ratio": 1.0    # No clipping
}

# Add quantization wrapper to each attention layer
for layer in model.model.layers:
    add_qk_rotation_wrapper_after_function_call_in_forward(
        layer.self_attn,       # Target module
        "apply_rotary_pos_emb",  # Hook after RoPE function
        config=model.config,   # Model configuration
        **k_quant_config       # Quantization parameters
    )

# The wrapper performs the following steps:
# 1. Applies Hadamard transform to query: q = Had(q) / sqrt(head_dim)
# 2. Applies Hadamard transform to key: k = Had(k) / sqrt(head_dim)
# 3. Quantizes the transformed key (k) with the configured bit width and groupsize.
```

### Parameters for K-cache Quantization

- **k_bits** (int) - Required - The number of bits to use for the key cache quantization.
- **k_groupsize** (int) - Required - The group size for quantization. Use 128 for standard grouping or -1 for token-wise quantization.
- **k_sym** (bool) - Required - Whether to use symmetric quantization. False for asymmetric.
- **k_clip_ratio** (float) - Required - Clipping ratio for quantization. 1.0 means no clipping.
```

--------------------------------

### Configure Activation Quantizer

Source: https://context7.com/facebookresearch/spinquant/llms.txt

The ActQuantizer class manages activation quantization, supporting symmetric or asymmetric modes and per-token or group-wise quantization strategies.

```python
from utils.quant_utils import ActQuantizer
quantizer = ActQuantizer()
quantizer.configure(bits=4, groupsize=-1, sym=False, clip_ratio=1.0)
x = torch.randn(1, 2048, 4096)
quantizer.find_params(x)
x_quantized = quantizer(x)
quantizer.free()
```

--------------------------------

### ActQuantizer Configuration

Source: https://context7.com/facebookresearch/spinquant/llms.txt

Handles activation quantization with support for symmetric/asymmetric modes and per-token or group-wise quantization.

```APIDOC
## ActQuantizer Configuration

The `ActQuantizer` class handles activation quantization with support for symmetric/asymmetric modes and per-token or group-wise quantization.

### Description
This class provides functionality to quantize activation tensors, allowing configuration of bit depth, quantization mode (symmetric/asymmetric), and grouping strategy (per-token or group-wise).

### Methods
- `configure(**kwargs)`: Configures the quantization parameters.
  - `bits` (int): Quantization bits (4, 8, or 16).
  - `groupsize` (int): -1 for per-token, positive for group-wise.
  - `sym` (bool): False for asymmetric, True for symmetric.
  - `clip_ratio` (float): Clipping ratio (0-1], 1.0 means no clipping.
- `find_params(x)`: Computes scale and zero point for quantization.
- `quantize(x)`: Applies quantization to the input tensor `x` using Straight-Through Estimator (STE).
- `free()`: Releases memory associated with the quantizer.

### Usage Example
```python
from utils.quant_utils import ActQuantizer

# Configure activation quantizer
quantizer = ActQuantizer()
quantizer.configure(
    bits=4,
    groupsize=-1,
    sym=False,
    clip_ratio=1.0
)

# Quantize activations during forward pass
x = torch.randn(1, 2048, 4096)  # (batch, seq_len, hidden_size)
quantizer.find_params(x)
x_quantized = quantizer(x)
quantizer.free()
```
```

--------------------------------

### Configure Weight Quantizer with MSE Clipping

Source: https://context7.com/facebookresearch/spinquant/llms.txt

The WeightQuantizer provides per-channel and group-wise quantization for model weights, featuring optional MSE-based clipping to minimize quantization error.

```python
from utils.quant_utils import WeightQuantizer
weight_quantizer = WeightQuantizer()
weight_quantizer.configure(bits=4, perchannel=True, sym=True, mse=True, norm=2.4, grid=100, maxshrink=0.8, weight_groupsize=128)
weight = model.layers[0].self_attn.q_proj.weight
weight_quantizer.find_params(weight)
quantized_weight = weight_quantizer.quantize(weight)
fake_quant, int_values, scales = weight_quantizer.fake_quantize(weight)
```

--------------------------------

### WeightQuantizer with MSE Optimization

Source: https://context7.com/facebookresearch/spinquant/llms.txt

Supports per-channel and group-wise weight quantization with optional MSE-based clipping optimization.

```APIDOC
## WeightQuantizer with MSE Optimization

The `WeightQuantizer` class supports per-channel and group-wise weight quantization with optional MSE-based clipping optimization.

### Description
This class enables quantization of weight tensors, offering flexibility in quantization strategy (per-channel, group-wise) and an optimization method (MSE) to determine clipping ranges.

### Methods
- `configure(**kwargs)`: Configures the weight quantization parameters.
  - `bits` (int): Target bit width.
  - `perchannel` (bool): True for per-channel quantization.
  - `sym` (bool): True for symmetric quantization.
  - `mse` (bool): True to enable MSE-based clipping search.
  - `norm` (float): Norm for MSE calculation.
  - `grid` (int): Grid search granularity.
  - `maxshrink` (float): Maximum shrink ratio for clipping.
  - `weight_groupsize` (int): Group size (-1 for per-tensor).
- `find_params(weight)`: Computes quantization parameters (scale, zero point, clipping) for the weight tensor.
- `quantize(weight)`: Applies quantization to the weight tensor.
- `fake_quantize(weight)`: Performs fake quantization and returns quantized values, integer values, and scales.

### Usage Example
```python
from utils.quant_utils import WeightQuantizer

# Configure weight quantizer with MSE optimization
weight_quantizer = WeightQuantizer()
weight_quantizer.configure(
    bits=4,
    perchannel=True,
    sym=True,
    mse=True,
    norm=2.4,
    grid=100,
    maxshrink=0.8,
    weight_groupsize=128
)

# Quantize weight tensor
weight = model.layers[0].self_attn.q_proj.weight
weight_quantizer.find_params(weight)
quantized_weight = weight_quantizer.quantize(weight)

# Get integer values and scales for export
fake_quant, int_values, scales = weight_quantizer.fake_quantize(weight)
```
```

--------------------------------

### Apply QKRotationWrapper for KV-Cache Quantization

Source: https://context7.com/facebookresearch/spinquant/llms.txt

Hooks into attention layers to apply Hadamard transforms and quantization to keys after RoPE. This is used to optimize KV-cache memory footprint.

```python
from eval_utils.rotation_utils import add_qk_rotation_wrapper_after_function_call_in_forward

k_quant_config = {"k_bits": 4, "k_groupsize": 128, "k_sym": False, "k_clip_ratio": 1.0}

for layer in model.model.layers:
    add_qk_rotation_wrapper_after_function_call_in_forward(
        layer.self_attn,
        "apply_rotary_pos_emb",
        config=model.config,
        **k_quant_config
    )
```

--------------------------------

### Evaluate Model Perplexity

Source: https://context7.com/facebookresearch/spinquant/llms.txt

Computes WikiText-2 perplexity using the evaluator utility. It processes the model layer-by-layer to optimize memory usage during evaluation.

```python
from utils.eval_utils import evaluator

class EvalArgs:
    bsz = 4
    capture_layer_io = False
    layer_idx = 0

args = EvalArgs()
model.seqlen = 2048

ppl = evaluator(model=model, testenc=testenc, dev="cuda", args=args)
print(f"WikiText-2 Perplexity: {ppl:.2f}")
```

--------------------------------

### Initialize RotateModule for Hidden State and Attention Head Rotations

Source: https://context7.com/facebookresearch/spinquant/llms.txt

The RotateModule wraps rotation matrices as trainable parameters. It is used to apply R1 rotations to hidden states and R2 rotations to individual attention heads.

```python
import torch
from torch import nn

class RotateModule(nn.Module):
    def __init__(self, R_init):
        super(RotateModule, self).__init__()
        self.weight = nn.Parameter(R_init.to(torch.float32).to(torch.device("cuda")))

    def forward(self, x, transpose=False):
        if transpose:
            return x @ self.weight
        else:
            return self.weight @ x

from utils.hadamard_utils import random_hadamard_matrix
hidden_size = 4096
head_dim = 128
R1 = random_hadamard_matrix(hidden_size, "cuda")
model.R1 = RotateModule(R1)
for i in range(num_layers):
    R2 = random_hadamard_matrix(head_dim, "cuda")
    model.model.layers[i].self_attn.R2 = RotateModule(R2)
```

--------------------------------

### RotateModule Class

Source: https://context7.com/facebookresearch/spinquant/llms.txt

Wraps learned rotation matrices as trainable parameters for hidden state and attention head rotations.

```APIDOC
## RotateModule Class

The `RotateModule` wraps learned rotation matrices as trainable parameters. It's used for both R1 (hidden state rotation) and R2 (attention head rotation).

### Description
This module initializes and manages trainable rotation matrices, applying them during the forward pass either directly or transposed.

### Initialization
- `R_init` (torch.Tensor): Initial rotation matrix.

### Forward Pass
- `x` (torch.Tensor): Input tensor.
- `transpose` (bool): If True, applies the transposed rotation matrix; otherwise, applies the original matrix.

### Usage Example
```python
import torch
from torch import nn

class RotateModule(nn.Module):
    def __init__(self, R_init):
        super(RotateModule, self).__init__()
        self.weight = nn.Parameter(R_init.to(torch.float32).to(torch.device("cuda")))

    def forward(self, x, transpose=False):
        if transpose:
            return x @ self.weight
        else:
            return self.weight @ x

# Initialize rotation modules for a model
from utils.hadamard_utils import random_hadamard_matrix

hidden_size = 4096  # LLaMA-2 7B
head_dim = 128      # hidden_size / num_heads

# R1 rotates the hidden states
R1 = random_hadamard_matrix(hidden_size, "cuda")
model.R1 = RotateModule(R1)

# R2 rotates each attention head independently
for i in range(num_layers):
    R2 = random_hadamard_matrix(head_dim, "cuda")
    model.model.layers[i].self_attn.R2 = RotateModule(R2)
```
```

--------------------------------

### Optimize Rotation Matrices with SGDG Optimizer

Source: https://context7.com/facebookresearch/spinquant/llms.txt

The SGDG optimizer performs gradient descent on the Stiefel manifold using the Cayley transform. This ensures that rotation matrices remain orthogonal throughout the training process.

```python
from train_utils.optimizer import SGDG
trainable_parameters = [model.R1.weight] + [model.model.layers[i].self_attn.R2.weight for i in range(model.config.num_hidden_layers)]
optimizer = SGDG(trainable_parameters, lr=1.5, momentum=0, stiefel=True)
for batch in dataloader:
    optimizer.zero_grad()
    loss = model(batch).loss
    loss.backward()
    optimizer.step()
```

--------------------------------

### Model Rotation Pipeline

Source: https://context7.com/facebookresearch/spinquant/llms.txt

Applies learned rotations to all model weights, fusing layer norms and applying Hadamard transforms.

```APIDOC
## Model Rotation Pipeline

The `rotate_model` function applies learned rotations to all model weights, fusing layer norms and applying Hadamard transforms.

### Description
This function integrates pre-trained rotation matrices into a model's weights, optionally fusing layer normalization operations and applying Hadamard transforms for efficiency and performance.

### Parameters
- `model` (nn.Module): The model to apply rotations to.
- `args` (object): An object containing configuration arguments for rotation.
  - `rotate_mode` (str): The mode of rotation to apply (e.g., "hadamard").
  - `optimized_rotation_path` (str): Path to the pre-trained rotation checkpoint.

### Steps Performed
1.  **Rotate embeddings**: `W_emb = W_emb @ R1`
2.  **Rotate attention**: `W_qkv = W_qkv @ R1`, `W_o = R1.T @ W_o`
3.  **Rotate MLP**: `W_up/gate = W_up/gate @ R1`, `W_down = R1.T @ W_down`
4.  **Apply Hadamard**: Apply `R2` to `V` and `O` projections.
5.  **Rotate lm_head**: `W_head = W_head @ R1`

### Usage Example
```python
from eval_utils.rotation_utils import rotate_model, get_orthogonal_matrix
from utils.fuse_norm_utils import fuse_layer_norms
import torch

# Load model and rotation checkpoint
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b")
rotation_checkpoint = torch.load("./rotations/llama2-7b-w4a4kv4/R.bin")

# Configure rotation arguments
class Args:
    rotate_mode = "hadamard"
    optimized_rotation_path = "./rotations/llama2-7b-w4a4kv4/R.bin"

args = Args()

# Apply rotations to model
model.eval()
fuse_layer_norms(model)  # Fuse RMSNorm into adjacent weights
rotate_model(model, args)
```
```

--------------------------------

### SGDG Optimizer (Cayley Transform)

Source: https://context7.com/facebookresearch/spinquant/llms.txt

Performs gradient descent on the Stiefel manifold using the Cayley transform to ensure rotation matrices remain orthogonal.

```APIDOC
## SGDG Optimizer (Cayley Transform)

The `SGDG` optimizer performs gradient descent on the Stiefel manifold using Cayley transform, ensuring rotation matrices remain orthogonal during training.

### Description
This optimizer is specifically designed for training parameters that must remain orthogonal, such as rotation matrices, by applying the Cayley transform to update them.

### Parameters
- `params` (iterable): An iterable of parameters to optimize.
- `lr` (float): Learning rate. Typically 1.0-2.0 for rotation optimization.
- `momentum` (float): Momentum factor (default: 0).
- `stiefel` (bool): If True, enables the Cayley transform for orthogonal optimization.

### Usage Example
```python
from train_utils.optimizer import SGDG

# Collect trainable rotation parameters
trainable_parameters = [model.R1.weight] + [
    model.model.layers[i].self_attn.R2.weight
    for i in range(model.config.num_hidden_layers)
]

# Initialize SGDG optimizer with Stiefel manifold constraint
optimizer = SGDG(
    trainable_parameters,
    lr=1.5,
    momentum=0,
    stiefel=True
)

# Training loop
for batch in dataloader:
    optimizer.zero_grad()
    loss = model(batch).loss
    loss.backward()
    optimizer.step()  # Updates rotations while maintaining orthogonality
```
```

--------------------------------

### Apply Model Rotation Pipeline

Source: https://context7.com/facebookresearch/spinquant/llms.txt

The rotate_model function applies learned rotation matrices to model weights, including embedding, attention, MLP, and lm_head layers, while fusing layer norms.

```python
from eval_utils.rotation_utils import rotate_model, get_orthogonal_matrix
from utils.fuse_norm_utils import fuse_layer_norms
import torch
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b")
class Args:
    rotate_mode = "hadamard"
    optimized_rotation_path = "./rotations/llama2-7b-w4a4kv4/R.bin"
args = Args()
model.eval()
fuse_layer_norms(model)
rotate_model(model, args)
```

--------------------------------

### Perplexity Evaluator

Source: https://context7.com/facebookresearch/spinquant/llms.txt

This section describes the `evaluator` function, which computes the perplexity of a model on the WikiText-2 dataset by processing the model layer-by-layer to optimize memory usage.

```APIDOC
## Perplexity Evaluator

The `evaluator` function computes WikiText-2 perplexity by processing the model layer-by-layer to minimize memory usage.

### Usage

```python
from utils.eval_utils import evaluator
from utils.data_utils import get_wikitext2
from transformers import AutoTokenizer

# Assume tokenizer and model are already initialized
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b")
# model = ... (your initialized model)

# Prepare evaluation data
testenc = get_wikitext2(seqlen=2048, tokenizer=tokenizer, eval_mode=True)

# Configure evaluation arguments
class EvalArgs:
    bsz = 4              # Batch size for evaluation
    capture_layer_io = False
    layer_idx = 0

args = EvalArgs()
model.seqlen = 2048

# Run evaluation
ppl = evaluator(
    model=model,
    testenc=testenc,
    dev="cuda",
    args=args
)

print(f"WikiText-2 Perplexity: {ppl:.2f}")
# Expected output for SpinQuant W4A4KV4 on LLaMA-2 7B: ~5.9
```

### Parameters for `evaluator`

- **model** (object) - Required - The model to evaluate.
- **testenc** (object) - Required - The tokenized test dataset (output of `get_wikitext2` with `eval_mode=True`).
- **dev** (str) - Required - The device to run evaluation on (e.g., "cuda").
- **args** (object) - Required - An object containing evaluation arguments (e.g., `bsz`, `capture_layer_io`, `layer_idx`).
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.