Optimize sparse 2:4 compression performance

[rahul-tuli]

Optimize Sparse 2:4 Compression Performance (2.64x Speedup)

Summary

This PR optimizes the sparse 2:4 compression pipeline, achieving a 2.64x speedup (30.18s → 11.42s) for Llama-3-8B sparse models by eliminating unnecessary CPU transfers and implementing vectorized GPU-accelerated bit packing operations.

Motivation

Profiling revealed that sparse compression was significantly slower than expected due to:

• Forced CPU execution via .cpu() calls before compression
• NumPy-based bit packing operations requiring tensor transfers
• Inefficient top-k selection parameters

Changes

1. Vectorized Bit Packing (src/compressed_tensors/utils/helpers.py)

• Implemented vectorized PyTorch-based pack_bitmasks() that efficiently packs bits using tensor operations
• Smart device handling: NumPy for CPU (optimal), PyTorch for GPU (avoids transfers)
• Refactored following DRY and single responsibility principles
• ~100x faster than loop-based implementation

# Before: Forces CPU transfer
packed_bits_numpy = numpy.packbits(bytemasks.numpy(), axis=-1, bitorder="little")

# After: Vectorized GPU implementation
bit_shifts = (1 << torch.arange(8, device=device, dtype=torch.uint8))
packed = (reshaped * bit_shifts).sum(dim=2, dtype=torch.uint8)

2. Deferred CPU Transfer (src/compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py)

• Removed premature .cpu() call in Sparse24BitMaskTensor.from_dense()
• Compression now occurs on the tensor's original device
• CPU transfer only happens at the end if needed for storage

3. Optimized Top-k Selection

• Added sorted=False parameter to topk() operation
• Sorting is unnecessary for 2:4 sparsity pattern selection
• ~10-15% performance improvement

Performance Results

Real Model Benchmark (Llama-3-8B-Instruct 2:4 Sparse)

Testing with /home/rahul/llm-compressor/Meta-Llama-3-8B-Instruct2of4-sparse on NVIDIA A100-80GB:

Branch	Compression Time	Speedup
main	30.18s	1.0x
optimized	11.42s	2.64x

Model details:

• Total parameters: 8.03B
• Sparse layers compressed: 224
• Compressed parameters: 739

Bit Packing Microbenchmark

Performance of the optimized pack_bitmasks function:

Size	CPU (NumPy)	GPU (PyTorch)	GPU Speedup
4096×4096	1.64ms	1.02ms	1.61x
8192×8192	5.72ms	1.74ms	3.29x

Verification

Quick Verification Script

# Save as verify_optimization.py
import torch
import numpy as np
from compressed_tensors.utils.helpers import pack_bitmasks

# Test correctness
mask = torch.rand(128, 256) > 0.5
packed_torch = pack_bitmasks(mask)
packed_numpy = torch.from_numpy(np.packbits(mask.numpy(), axis=-1, bitorder="little"))
assert torch.equal(packed_torch, packed_numpy), "Implementation mismatch!"

# Test GPU performance
if torch.cuda.is_available():
    mask_gpu = torch.rand(4096, 4096).cuda() > 0.5
    packed_gpu = pack_bitmasks(mask_gpu)
    assert packed_gpu.is_cuda, "Should stay on GPU"
    print("✓ GPU optimization working correctly")

Benchmark Script

import time
import torch
from transformers import AutoModelForCausalLM
from transformers.utils.quantization_config import CompressedTensorsConfig
from compressed_tensors import ModelCompressor, Sparse24BitMaskConfig

MODEL_PATH = "/home/rahul/llm-compressor/Meta-Llama-3-8B-Instruct2of4-sparse"

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype="auto",
    device_map="auto",
    quantization_config=CompressedTensorsConfig(run_compressed=False)
)

sparsity_config = Sparse24BitMaskConfig(
    format="sparse-24-bitmask",
    targets=['Linear'],
    ignore=['lm_head'],
)
compressor = ModelCompressor.from_pretrained_model(
    model,
    sparsity_config=sparsity_config,
    quantization_format=None
)

torch.cuda.synchronize()
start = time.time()
compressed_state_dict = compressor.compress(model)
torch.cuda.synchronize()
compression_time = time.time() - start

print(f"Compression time: {compression_time:.2f}s")

Run Tests

# Run the new test suite
pytest tests/test_sparse_optimization.py -v

# Run existing tests to ensure compatibility
pytest tests/test_compressors/test_sparse_compressors.py -v

Technical Details

Vectorized Algorithm

The new implementation uses vectorized operations instead of Python loops:

1. Reshapes tensor to process 8 bits at a time
2. Uses broadcasting with bit shift values [1, 2, 4, 8, 16, 32, 64, 128]
3. Performs element-wise multiplication and sum reduction
4. Results in ~100x speedup over loop-based approach

Memory Efficiency

• No increase in memory usage
• Operations remain in-place where possible
• Temporary buffers are minimal and device-local

Backward Compatibility

• All changes maintain backward compatibility
• NumPy fallback ensures existing CPU workflows continue to work
• No API changes

Code Quality

Refactoring

• Split functions following single responsibility principle
• Added comprehensive documentation
• Improved error messages with context
• Cleaner separation of concerns

Test Coverage

• 10 focused test cases
• Edge case handling (non-multiple-of-8, empty tensors, etc.)
• Performance regression tests
• Multiple dtype support

Checklist

• Code follows project style guidelines
• All tests pass
• Performance benchmarks included
• Verification scripts provided
• Backward compatibility maintained
• Error handling implemented
• Code refactored following best practices
• Real model benchmarks demonstrate 2.64x speedup