#1920 activation sparsity + compression

[ved1beta]

Activation Compression Module for PyTorch AO

Overview

This PR introduces a new activation compression module for PyTorch AO that enables efficient training of large neural networks by compressing sparse activations. The module provides significant memory savings while maintaining model accuracy.

Key Features

• Multiple compression methods:
  • simple: Basic compression storing non-zero values and indices
  • block: Block-based compression for better cache efficiency
  • run_length: Run-length encoding for sequences of zeros
• Memory-efficient training with automatic compression/decompression
• Comprehensive error handling and validation
• Memory usage tracking and statistics
• Seamless integration with existing PyTorch models

How It Works

Core Components

1. ActivationCompressor

   • Handles the compression and decompression of sparse tensors
   • Supports multiple compression methods
   • Tracks memory usage
   • Validates input tensors and compression methods

2. CompressedActivation

   • PyTorch module that manages activation compression during training
   • Automatically compresses activations in forward pass
   • Handles proper gradient computation in backward pass
   • Maintains sparsity patterns in gradients

Usage Example

import torch
from ao.sparsity import CompressedActivation

# Create a compressed activation layer
compressed_layer = CompressedActivation(compression_method='block')

# Use in your model
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(512, 2048)
        self.compressed = CompressedActivation()
        self.fc2 = nn.Linear(2048, 512)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.compressed(x)  # Activations are automatically compressed
        x = self.fc2(x)
        return x

Compression Methods

1. Simple Compression

   • Stores non-zero values and their indices
   • Best for general use cases
   • Minimal overhead

2. Block Compression

   • Groups values into blocks for better cache efficiency
   • Ideal for structured sparsity patterns
   • Configurable block size

3. Run-Length Encoding

   • Efficient for sequences with many zeros
   • Best for highly sparse activations
   • Minimal memory overhead

Performance Benefits

• Significant memory savings during training
• Minimal impact on model accuracy
• Efficient compression/decompression operations
• Automatic sparsity-aware gradient computation

Testing

The module includes comprehensive tests for:

• Compression/decompression correctness
• Memory usage tracking
• Gradient computation
• Different compression methods
• Edge cases and error handling

Documentation

• Detailed docstrings for all classes and methods
• Usage examples and best practices
• Performance benchmarks
• Integration guidelines

Dependencies

• PyTorch >= 1.8.0
• No additional dependencies required

Checklist

• Code follows the project's style guidelines
• Tests have been added/updated
• Documentation has been updated
• Performance benchmarks included
• No breaking changes introduced

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2076

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 8 New Failures

As of commit 93785f4 with merge base a61e302 ():

NEW FAILURES - The following jobs have failed:

• Code Analysis with Ruff / build (3.9) (gh)
  ao/sparsity/tests/test_activation_compression.py:67:9: F841 Local variable grad_input is assigned to but never used
• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.
• Run Regression Tests / test (CPU 2.4, linux.4xlarge, torch==2.4.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2
• Run Regression Tests / test (CPU 2.5.1, linux.4xlarge, torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2
• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2
• Run Regression Tests / test (CUDA 2.5.1, linux.g5.12xlarge.nvidia.gpu, torch==2.5.1 --index-url https://download.pytorch... / linux-job (gh)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[supriyar]

cc @jcaip