read/IMPLEMENTATION_GUIDE.md

# Implementation Guide

This document provides a roadmap for implementing the BitLinear functionality, following the structure defined in the project skeleton. This is here to give insight on how one can replicate this process to different operations. 

## Implementation Order

### Phase 1: Python Baseline (Correctness First)

Start here to establish correctness before optimizing.

#### 1.1 Quantization (`bitlinear/quantization.py`)

Order of implementation:
1. `absmax_scale()` - Simple max computation
2. `ternary_quantize()` - Threshold-based quantization to {-1, 0, +1}
3. `weight_to_ternary()` - Combines the above
4. Test thoroughly with `tests/test_quantization.py`

**Key considerations:**
- Threshold selection (try 0.33 * scale or 0.5 * scale)
- Per-channel vs. global scaling trade-offs
- Numerical stability (avoid division by zero)

#### 1.2 Functional Operations (`bitlinear/functional.py`)

Order of implementation:
1. `bitlinear_python()` - Core ternary matmul
   ```python

   # Pseudocode:

   output = torch.matmul(x, W_ternary.T)

   output = output * gamma.unsqueeze(0)

   if bias is not None:

       output = output + bias

   return output

   ```

2. `greedy_ternary_decomposition()` - Iterative residual quantization
   ```python

   # Pseudocode:

   residual = W.clone()

   for i in range(k):

       W_t, gamma = weight_to_ternary(residual)

       store W_t and gamma

       residual = residual - gamma * W_t

   ```

3. `multi_ternary_linear_python()` - Sum of k ternary operations

4. Test with `tests/test_functional.py`

#### 1.3 Layer Modules (`bitlinear/layers.py`)

Order of implementation:
1. `BitLinear.__init__()` and `reset_parameters()`
   - Initialize dense weights using kaiming_uniform

   - Quantize to ternary using `weight_to_ternary()`

   - Store as buffers or parameters



2. `BitLinear.forward()` - Call `bitlinear_python()`

3. `BitLinear.from_linear()` - Conversion utility

4. `MultiTernaryLinear` - Similar structure

5. `convert_linear_to_bitlinear()` - Recursive module conversion

6. Test with `tests/test_layers.py`

**Testing strategy:**
- Compare output shapes with nn.Linear
- Verify ternary weight values
- Test conversion from pre-trained weights
- Validate in Transformer example

### Phase 2: Memory Optimization

#### 2.1 Base-3 Packing (`bitlinear/packing.py`)

Implement packing for memory efficiency:
1. `pack_ternary_base3()` - 5 values per byte
2. `unpack_ternary_base3()` - Reverse operation
3. Verify roundtrip: pack → unpack == identity

**Packing scheme:**
```

Map: -1 → 0, 0 → 1, +1 → 2 (base-3 digits)

Pack 5 digits per byte: d0 + d1*3 + d2*9 + d3*27 + d4*81

```

### Phase 3: C++ Extensions (Optional but Recommended)

#### 3.1 CPU Implementation (`bitlinear/cpp/bitlinear.cpp`)

1. Implement `bitlinear_cpu_forward()`
   - Basic matrix multiplication with ternary weights
   - Exploit ternary structure (skip multiplications)

2. Implement `multi_ternary_cpu_forward()`

3. Test integration with Python

**Optimization opportunities (later):**
- AVX/AVX512 vectorization
- OpenMP parallelization
- Cache-efficient tiling

#### 3.2 CUDA Kernels (`bitlinear/cpp/bitlinear_kernel.cu`)



Only after CPU version works!



1. Basic kernel without optimization

   - Thread per output element

   - Simple accumulation



2. Optimized kernel:

   - Shared memory tiling

   - Warp-level reductions

   - Memory coalescing

   - Exploit ternary (conditional accumulation)



3. Advanced (optional):

   - Tensor Core utilization

   - Mixed precision

   - Fused kernels (activation quantization + matmul)



**Performance targets:**

- Should be faster than PyTorch's F.linear for large matrices

- Aim for 2-5x speedup from ternary optimization



### Phase 4: Training Support



#### 4.1 Quantization-Aware Training (QAT)



Modify layers to support gradient flow:

1. Straight-through estimator for ternary quantization

2. Learnable scaling factors (gamma)

3. Fine-tuning pre-trained models



#### 4.2 Initialization Strategies



Experiment with initialization for ternary weights:

- Standard kaiming_uniform then quantize
- Specialized initialization for ternary
- Better threshold selection

## Testing Strategy

### Unit Tests
Run frequently during development:
```bash

pytest tests/test_quantization.py -v

pytest tests/test_functional.py -v

pytest tests/test_layers.py -v

```

### Integration Tests
Test full pipelines:
1. Dense model → quantization → inference
2. Transformer with BitLinear layers
3. Save/load model checkpoints

### Numerical Correctness
Compare with reference:
```python

# Create same layer in dense and ternary

linear = nn.Linear(512, 512)

bitlinear = BitLinear.from_linear(linear)



x = torch.randn(32, 512)

out_dense = linear(x)

out_ternary = bitlinear(x)



# Should be similar (not identical due to quantization)

error = torch.norm(out_dense - out_ternary) / torch.norm(out_dense)

print(f"Relative error: {error:.4f}")  # Expect ~0.1-0.3

```

## Common Pitfalls

### Quantization
- **Pitfall:** Wrong threshold → too many zeros or not enough
- **Solution:** Start with 0.5 * scale, tune empirically

### Shape Handling
- **Pitfall:** Broadcasting errors with gamma
- **Solution:** Use `.unsqueeze()` carefully, test various input shapes

### CUDA Compilation
- **Pitfall:** CUDA version mismatches
- **Solution:** Match PyTorch's CUDA version, use CPU-only build first

### Gradients
- **Pitfall:** No gradient flow through ternary quantization
- **Solution:** Implement straight-through estimator for QAT

## Performance Benchmarks

Create benchmarks to track progress:
```python

import time

import torch

from bitlinear import BitLinear



def benchmark(layer, x, n_runs=100):

    # Warmup

    for _ in range(10):

        _ = layer(x)

    

    # Benchmark

    start = time.time()

    for _ in range(n_runs):

        _ = layer(x)

    end = time.time()

    

    return (end - start) / n_runs



# Compare

linear = nn.Linear(2048, 2048).cuda()

bitlinear = BitLinear(2048, 2048).cuda()

x = torch.randn(128, 2048).cuda()



time_linear = benchmark(linear, x)

time_bitlinear = benchmark(bitlinear, x)



print(f"nn.Linear: {time_linear*1000:.2f} ms")

print(f"BitLinear: {time_bitlinear*1000:.2f} ms")

print(f"Speedup: {time_linear/time_bitlinear:.2f}x")

```

## Next Steps After Skeleton

1. **Implement Phase 1** (Python baseline)
   - Start with `absmax_scale()` and `ternary_quantize()`
   - Test each function as you go
   - Don't move to next phase until tests pass

2. **Validate with Examples**
   - Run `examples/basic_usage.py`
   - Run `examples/transformer_example.py`
   - Check output similarity and memory savings

3. **Optimize if Needed**
   - Profile to find bottlenecks
   - Implement C++/CUDA only after Python works
   - Measure performance improvements

4. **Documentation**
   - Add docstring details from implementation
   - Create API documentation
   - Write usage tutorials

## Resources

### Papers
- BitNet: https://arxiv.org/abs/2310.11453
- Ternary Neural Networks: https://jmlr.org/papers/volume26/24-2050/24-2050.pdf

### PyTorch Resources
- Custom Extensions: https://pytorch.org/tutorials/advanced/cpp_extension.html

- CUDA Programming: https://pytorch.org/tutorials/advanced/custom_ops.html

### Quantization
- QAT Guide: https://pytorch.org/docs/stable/quantization.html
- Straight-through Estimator: Bengio et al., 2013

## Questions to Consider

As you implement, think about:
1. **Memory vs. Speed:** Packed weights save memory but need unpacking
2. **Training vs. Inference:** Different requirements for gradients
3. **Compatibility:** Should work with existing PyTorch features (DDP, AMP, etc.)
4. **Extensibility:** Easy to add new quantization schemes?

Good luck with implementation! Start with correctness, then optimize.