torch-floating-point 0.0.6

A PyTorch library for custom floating point quantization with autograd support

Torch Floating Point

A PyTorch library for custom floating point quantization with autograd support. This library provides efficient implementations of custom floating point formats with automatic differentiation capabilities.

Features

• Custom Floating Point Formats: Support for arbitrary floating point configurations (sign bits, exponent bits, mantissa bits, bias)
• Autograd Support: Full PyTorch autograd integration for training with quantized weights
• CUDA Support: GPU acceleration for both forward and backward passes
• Multiple Precision: Support for various bit widths (4-bit, 8-bit, 16-bit, 32-bit)
• Straight-Through Estimator: Gradient-friendly quantization for training
• Comprehensive Testing: Extensive test suite covering differentiability and accuracy

Installation

From PyPI (Recommended)

pip install torch-floating-point

From Source

git clone https://github.com/SamirMoustafa/torch-floating-point.git
cd torch-floating-point
pip install -e .

Development Installation

git clone https://github.com/SamirMoustafa/torch-floating-point.git
cd torch-floating-point
pip install -e ".[dev,test]"
pre-commit install

Quick Start

import torch
from floating_point import FloatingPoint, Round

# Define a custom 8-bit floating point format (1 sign, 4 exponent, 3 mantissa bits)
fp8 = FloatingPoint(sign_bits=1, exponent_bits=4, mantissa_bits=3, bias=7, bits=8)

# Create a rounding function
rounder = Round(fp8)

# Create a tensor with gradients
x = torch.randn(10, requires_grad=True)

# Quantize the tensor
quantized = rounder(x)

# Use in training (gradients flow through)
loss = quantized.sum()
loss.backward()

print(f"Original: {x}")
print(f"Quantized: {quantized}")
print(f"Gradients: {x.grad}")

Examples

The project includes comprehensive examples in the examples/ directory:

Simple Rounding Example (examples/01_simple_rounding.py)

Demonstrates basic rounding functionality with different floating point formats:

• Compares FP4, FP8, and FP16 precision
• Shows quantization errors and differences
• Demonstrates range limitations and edge cases
• Handles subnormal values

Gradient Flow Example (examples/02_gradient_flow.py)

Demonstrates gradient flow through quantized operations:

• Shows Straight-Through Estimator (STE) in action
• Compares gradient flow across different formats
• Includes a complete training loop with quantized weights
• Analyzes gradient patterns and clipping behavior

Running Examples

# Run individual examples
python examples/01_simple_rounding.py
python examples/02_gradient_flow.py

# Run all examples
python examples/run_all_examples.py

Usage Examples

Custom Floating Point Configuration

from floating_point import FloatingPoint

# 4-bit floating point (1 sign, 2 exponent, 1 mantissa)
fp4 = FloatingPoint(sign_bits=1, exponent_bits=2, mantissa_bits=1, bias=1, bits=4)

# 8-bit floating point with custom max mantissa
fp8_custom = FloatingPoint(
    sign_bits=1, 
    exponent_bits=4, 
    mantissa_bits=3, 
    bias=7, 
    bits=8,
    max_mantissa_at_max_exponent=6,  # Custom max mantissa
    reserved_exponent=False  # No reserved exponent for inf/nan
)

# 16-bit floating point (standard)
fp16 = FloatingPoint(sign_bits=1, exponent_bits=5, mantissa_bits=10, bias=15, bits=16)

Training with Quantized Weights

import torch
import torch.nn as nn
from floating_point import FloatingPoint, Round

class QuantizedLinear(nn.Module):
    def __init__(self, in_features, out_features, fp_config):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features))
        self.rounder = Round(fp_config)
    
    def forward(self, x):
        quantized_weight = self.rounder(self.weight)
        return torch.nn.functional.linear(x, quantized_weight)

# Define quantization format
fp8 = FloatingPoint(sign_bits=1, exponent_bits=4, mantissa_bits=3, bias=7, bits=8)

# Create model with quantized weights
model = QuantizedLinear(784, 10, fp8)
optimizer = torch.optim.Adam(model.parameters())

# Training loop
for epoch in range(10):
    # ... your training code ...
    loss.backward()
    optimizer.step()

Direct Function Usage

import torch
from floating_point import autograd

# Direct quantization function
x = torch.randn(100, requires_grad=True)
quantized = autograd(x, exponent_bits=4, mantissa_bits=3, bias=7)

# Gradients work automatically
loss = quantized.sum()
loss.backward()

Supported Formats

The library supports various floating point formats:

Format	Sign Bits	Exponent Bits	Mantissa Bits	Bias	Total Bits
FP4	1	2	1	1	4
FP8	1	4	3	7	8
FP16	1	5	10	15	16
BF16	1	8	7	127	16
FP32	1	8	23	127	32

Development

Testing

The project includes two testing approaches:

1. CI/CD Tests (GitHub Actions): Fast, lightweight tests that verify core functionality without heavy numerical computations
2. Full Test Suite: Complete test coverage including all numerical precision tests (run locally or via manual workflow trigger)

To run the full test suite locally:

export LD_LIBRARY_PATH=$(python -c "import torch; print(torch.__file__)")/lib:$LD_LIBRARY_PATH
python -m pytest test/round.py test/data_types.py -v

Running Tests

# Run all tests
make test

# Run tests with coverage
make test-cov

# Run specific test file
python -m pytest test/round.py -v

Code Quality

# Run linting
make lint

# Format code
make format

# Run all checks
make full-check

Building

# Build extension
python setup.py build_ext --inplace

# Build package
make build

# Clean build artifacts
make clean

Contributing

1. Fork the repository
2. Create a feature branch (git checkout -b feature/amazing-feature)
3. Install development dependencies (make setup-dev)
4. Make your changes
5. Run tests (make test)
6. Run linting (make lint)
7. Commit your changes (git commit -m 'Add amazing feature')
8. Push to the branch (git push origin feature/amazing-feature)
9. Open a Pull Request

License

This project is licensed under the MIT License - see the LICENSE file for details.

Citation

If you use this library in your research, please cite:

@software{torch_floating_point,
  title={Torch Floating Point: A PyTorch library for custom floating point quantization},
  author={Samir Moustafa},
  year={2024},
  url={https://github.com/SamirMoustafa/torch-floating-point}
}

Acknowledgments

• PyTorch team for the excellent autograd system
• The PyTorch C++ extension community for guidance on extension development
• Contributors and users of this library

Support

• Issues: GitHub Issues
• Discussions: GitHub Discussions
• Email: [e-mail]