GRACE-SVD

GRACE-SVD is a tool for equalizing cores in SVD decomposed models. It provides multiple equalization strategies that can be easily switched between, making it suitable for model compression and quantization workflows.

Features

• Multiple Equalization Strategies:

  • Whitening Division Factor: Divide whitening matrix by a given factor
  • Max-Min Range: Max-min range based equalization
  • Wasserstein Distance: Minimize Wasserstein distance using kurtosis-based range ratio

• Simple CLI: Run directly without installation

• Comprehensive Validation: Built-in validation for state dictionaries

• Analysis Tools: Compute and analyze scaling factors

Quick Start (No Installation Required)

Simply run the script directly:

# Navigate to the GRACE-SVD directory
cd /path/to/GRACE-SVD

# List available strategies
python core_equalizer.py --list-strategies

# Show information about a model
python core_equalizer.py --info /path/to/model.pt

# Equalize a model
python core_equalizer.py -i model.pt -o equalized.pt -s wasserstein_distance --px 0.5 --beta 1e-3

# Compute scaling factors (analysis)
python core_equalizer.py --compute-sf model.pt --px 0.5 --beta 1e-3

Dependencies

GRACE-SVD requires the following packages:

• torch >= 1.9.0
• numpy >= 1.20.0
• scipy >= 1.7.0

Install them with:

pip install -r requirements.txt

Or install manually:

pip install torch numpy scipy

Optional (for Excel export):

pip install pandas openpyxl

Usage Examples

Command Line Interface

# Equalize using Wasserstein distance strategy (default)
python core_equalizer.py --input model.pt --output equalized.pt --strategy wasserstein_distance --px 0.5 --beta 1e-3

# Equalize using max-min range strategy
python core_equalizer.py -i model.pt -o equalized.pt -s max_min_range

# Equalize using whitening division factor strategy with custom factor
python core_equalizer.py -i model.pt -o equalized.pt -s whitening_division_factor --factor 2.0

# Show information about a model
python core_equalizer.py --info model.pt

# List available strategies
python core_equalizer.py --list-strategies

# Compute scaling factors and save to Excel
python core_equalizer.py --compute-sf model.pt --sf-output scaling_factors.xlsx

Python API

# Run from GRACE-SVD directory
import sys
sys.path.insert(0, '/path/to/GRACE-SVD')
from core_equalizer import CoreEqualizer, EqualizationStrategy, equalize_cores

# Create equalizer with Wasserstein distance strategy (default)
equalizer = CoreEqualizer(strategy=EqualizationStrategy.WASSERSTEIN_DISTANCE)

# Equalize a model file
equalizer.equalize_file(
    input_path="unequalized_model.pt",
    output_path="equalized_model.pt",
    px=0.5,           # Percentile limit
    beta=1e-3         # Kurtosis scaling factor
)

# Or use the convenience function
from equalizer import equalize_cores

equalize_cores(
    "model.pt",
    "equalized.pt",
    strategy="wasserstein_distance",
    px=0.5,
    beta=1e-3
)

Computing Scaling Factors (Analysis Tool)

The --compute-sf option allows you to analyze and compute scaling factors for all node pairs in your model before equalization. This is useful for understanding the distribution of scaling factors and debugging.

Usage

# Basic usage - compute and display scaling factors
python core_equalizer.py --compute-sf model.pt

# With custom parameters
python core_equalizer.py --compute-sf model.pt --px 0.5 --beta 1e-3

# Save results to Excel for further analysis
python core_equalizer.py --compute-sf model.pt --sf-output scaling_factors.xlsx

Output

The script outputs:

1. Scaling Factors Table: Shows SF (MaxMin), SF (Wasserstein), and sqrt(SF) for each layer
2. Summary Statistics: Average, min, max, and std of scaling factors
3. Range Statistics: Average ranges for Node0 and Node1

Equalization Strategies

1. Whitening Division Factor Strategy

Divides the whitening matrix by a given factor. Useful when you have pre-computed whitening matrices from calibration data.

Parameters:

• whitening_matrix_division_factor (float): Pre-computed scaling factor (default: 1.0)

python core_equalizer.py -i model.pt -o equalized.pt -s whitening_division_factor --factor 1.0

2. Max-Min Range Strategy

Uses the max-min range of weight values to compute scaling factors. The scaling factor is based on the full range of weight values.

Parameters: None required

python core_equalizer.py -i model.pt -o equalized.pt -s max_min_range

3. Wasserstein Distance Strategy (Default)

Minimizes the Wasserstein distance between node distributions by using kurtosis-based adaptive percentile range to compute scaling factors.

Parameters:

• px (float): Base percentile limit (default: 0.5)
• beta (float): Scaling factor for kurtosis influence (default: 1e-3)

Formula:

adaptive_px = px * tanh(beta * kurtosis)

python core_equalizer.py -i model.pt -o equalized.pt -s wasserstein_distance --px 0.5 --beta 1e-3

How It Works

Required State Dictionary Format

IMPORTANT: Your .pt file must contain tensors with nodes.0 and nodes.1 extensions. GRACE-SVD identifies node pairs by looking for keys matching this pattern:

# Example valid tensor names:
model.layers.0.self_attn.q_proj.nodes.0   <-- Node 0
model.layers.0.self_attn.q_proj.nodes.1   <-- Node 1 (paired with nodes.0)

model.layers.5.mlp.gate_proj.nodes.0      <-- Node 0
model.layers.5.mlp.gate_proj.nodes.1      <-- Node 1 (paired with nodes.0)

The Equalization Process

1. Identify Node Pairs: Find all pairs of cores (nodes.X.0 and nodes.X.1) in the state dictionary
2. Compute Ranges: Calculate the value range for each core using the selected strategy
3. Calculate Scaling Factor: SF = range_node1 / range_node0
4. Apply Scaling:
   • Node 0 (nodes.0): MULTIPLIED by sqrt(SF)
   • Node 1 (nodes.1): DIVIDED by sqrt(SF)

How Scaling Factors Are Applied to Nodes

Scaling Factor (SF) = Range_Node1 / Range_Node0

Node0 (nodes.0):  nodes.0 = nodes.0 * sqrt(SF)    [MULTIPLIED by sqrt(SF)]
Node1 (nodes.1):  nodes.1 = nodes.1 / sqrt(SF)    [DIVIDED by sqrt(SF)]

Why sqrt(SF)?

Since we apply scaling to both nodes, using sqrt ensures that:

• Node0's range increases by sqrt(SF)
• Node1's range decreases by sqrt(SF)
• Combined effect: both ranges become approximately equal

Example

If Node0 has range 0.01 and Node1 has range 0.1:

SF = 0.1 / 0.01 = 10
sqrt(SF) = 3.162

After equalization:
  Node0 range ≈ 0.01 * 3.162 = 0.03162
  Node1 range ≈ 0.1 / 3.162 = 0.03162

Both nodes now have equal ranges!

This process ensures that both cores have similar value ranges, which is important for:

• Quantization accuracy
• Numerical stability
• Model compression quality

API Reference

CoreEqualizer Class

class CoreEqualizer:
    def __init__(self, strategy: Union[EqualizationStrategy, str] = "wasserstein_distance"):
        """Initialize with specified strategy."""
    
    def equalize(
        self,
        state_dict: Dict[str, torch.Tensor],
        validate: bool = True,
        verbose: bool = True,
        **kwargs
    ) -> Dict[str, torch.Tensor]:
        """Equalize cores in the given state dictionary."""
    
    def equalize_file(
        self,
        input_path: Union[str, Path],
        output_path: Union[str, Path],
        validate: bool = True,
        verbose: bool = True,
        **kwargs
    ) -> Dict[str, torch.Tensor]:
        """Load, equalize, and save a state dictionary."""
    
    def set_strategy(self, strategy: Union[EqualizationStrategy, str]) -> None:
        """Change the equalization strategy."""
    
    @staticmethod
    def list_strategies() -> Dict[str, str]:
        """List all available strategies with descriptions."""

Convenience Function

def equalize_cores(
    input_path: Union[str, Path],
    output_path: Union[str, Path],
    strategy: Union[EqualizationStrategy, str] = "wasserstein_distance",
    **kwargs
) -> Dict[str, torch.Tensor]:
    """Quick equalization without creating a CoreEqualizer instance."""

Utility Functions

from utils import (
    load_state_dict,           # Load .pt file
    save_state_dict,           # Save .pt file
    get_state_dict_info,       # Get model info
    print_state_dict_info,     # Print model info
    validate_state_dict_for_equalization,  # Validate for equalization
    extract_layer_info,        # Extract layer metadata
)