Bitpacking (pytorch#291)

test/prototype/test_bitpacking.py

@@ -0,0 +1,70 @@
+import torch
+from torchao.prototype.common.bitpacking import pack, unpack
+import pytest
+from torch.utils._triton import has_triton
+from torchao.quantization.utils import TORCH_VERSION_AFTER_2_4
+
+if not TORCH_VERSION_AFTER_2_4:
+    pytest.skip("Unsupported PyTorch version", allow_module_level=True)
+
+def test_uint4_to_uint8_CPU():
+    test_tensor = torch.randint(0, 15, (4, 4), dtype=torch.uint8)
+    packed = pack(test_tensor, 8, 4, device='cpu')
+    unpacked = unpack(packed, 4, device='cpu')
+    unpadded = unpacked[:test_tensor.shape[0], ...]
+    assert(unpadded.allclose(test_tensor))
+
+def test_uint3_to_int16_col_wise_cpu():
+    test_tensor = torch.randint(0, 7, (8, 5), dtype=torch.int16)
+    packed = pack(test_tensor,16, 3, False, device='cpu')
+    unpacked = unpack(packed, 3, False, device='cpu')
+    unpadded = unpacked[:test_tensor.shape[0], ...]
+    assert(unpadded.allclose(test_tensor))
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_uint4_to_uint8():
+    test_tensor = torch.randint(0, 15, (4, 4), dtype=torch.uint8).cuda()
+    packed = pack(test_tensor, 8, 4)
+    unpacked = unpack(packed, 4)
+    unpadded = unpacked[:test_tensor.shape[0], ...]
+    assert(unpadded.allclose(test_tensor))
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.skipif(not has_triton(), reason="unsupported without triton")
+def test_uint4_to_uint8_compile():
+    torch._dynamo.config.specialize_int = True
+    pack_compiled = torch.compile(pack, fullgraph=True)
+    unpack_compiled = torch.compile(unpack, fullgraph=True)
+    test_tensor = torch.randint(0, 15, (3, 4), dtype=torch.uint8).cuda()
+    packed = pack_compiled(test_tensor, 8, 4)
+    unpacked = unpack_compiled(packed, 4)
+    unpadded = unpacked[:test_tensor.shape[0], ...]
+    assert(unpadded.allclose(test_tensor))
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_uint3_to_int16():
+    test_tensor = torch.randint(0, 7, (5, 8), dtype=torch.int16).cuda()
+    packed = pack(test_tensor,16, 3)
+    unpacked = unpack(packed, 3)
+    unpadded = unpacked[:test_tensor.shape[0], ...]
+    assert(unpadded.allclose(test_tensor))
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.skipif(not has_triton(), reason="unsupported without triton")
+def test_uint2_to_uint8_col_wise_compile():
+    torch._dynamo.config.specialize_int = True
+    pack_compiled = torch.compile(pack, fullgraph=True)
+    unpack_compiled = torch.compile(unpack, fullgraph=True)
+    test_tensor = torch.randint(0, 3, (8, 8), dtype=torch.uint8).cuda()
+    packed = pack_compiled(test_tensor, 8, 2, False)
+    unpacked = unpack_compiled(packed,2, False)
+    unpadded = unpacked[:test_tensor.shape[0], ...]
+    assert(unpadded.allclose(test_tensor))
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_uint3_to_int16_col_wise():
+    test_tensor = torch.randint(0, 7, (8, 5), dtype=torch.int16).cuda()
+    packed = pack(test_tensor,16, 3, False)
+    unpacked = unpack(packed, 3, False)
+    unpadded = unpacked[:test_tensor.shape[0], ...]
+    assert(unpadded.allclose(test_tensor))

torchao/prototype/common/bitpacking.py

@@ -0,0 +1,101 @@
+import torch
+from functools import reduce
+
+
+
+def unpack(data, data_size, by_rows = True, device="cuda"):
+    """
+    Unpacks small dtype elements from a larger dtype.
+    
+    Inputs:
+    data: torch.Tensor - a tensor of packed elements of a small dtype within a larger dtype.
+    data_size: int - the size of the small dtype in bits.
+    
+    optional:
+    by_rows: bool - specifies whether to unpack... 
+        by rows: tensor(n,m) -> tensor(n*scale, m) 
+        or by columns: tensor(n,m) -> tensor(n,m*scale)
+        
+    defaults to rows because quantization is typically done by rows 
+    but choose the version which matches how you quantize as this improves memory accesses/performance
+    
+    Returns: torch.Tensor - a tensor of the unpacked elements.
+    """
+    if by_rows:
+        return _unpack_by_rows(data, data_size, device)
+    else:
+        return _unpack_by_cols(data, data_size)
+
+def pack(data, container_size, data_size, by_rows = True, device="cuda"):
+    """
+    Packs small dtype elements into a larger dtype.
+    Pads rows to be divisible by the scale.
+    
+    Inputs:
+    data: torch.Tensor - a tensor of unpacked elements of a small dtype.
+    container_size: int - the size of the large dtype in bits.
+    data_size: int - the size of the small dtype in bits.
+    
+    optional:
+    by_rows: bool - specifies whether to pack values... 
+        by rows: tensor(n,m) -> tensor(n//scale, m) 
+        or by columns: tensor(n,m) -> tensor(n,m//scale)
+    
+    defaults to rows because quantization is typically done by rows
+    but choose the version which matches how you quantize as this improves memory accesses/performance
+    
+    Returns: torch.Tensor - a tensor of packed elements.
+    """
+    if by_rows:
+        return _pack_by_rows(data, container_size, data_size, device)
+    else:
+        return _pack_by_cols(data, container_size, data_size, device)   
+
+def _unpack_by_rows(data, data_size, device) -> torch.Tensor:
+    shape = data.shape
+    scale = data.element_size() * 8 // data_size
+
+    unpacked_data = torch.zeros((shape[0]*scale, *shape[1:]), dtype=data.dtype).to(device)
+    nbits = (1 << data_size) - 1 # mask for the last dtype_size bits
+    for i in range(scale):
+        shift_amt = data.element_size() * 8 - data_size * (i + 1) # how much to shift to get the ith uint
+        unpacked_data[i::scale] = ((data >> shift_amt) & (nbits))
+    return unpacked_data
+
+def _unpack_by_cols(data, data_size) -> torch.Tensor:
+    shape = data.shape
+    scale = data.element_size() * 8 // data_size
+    unpacked_data = []
+    nbits = (1 << data_size) - 1 # mask for the last dtype_size bits
+    for i in range(scale):
+        shift_amt = data.element_size() * 8 - data_size * (i + 1) # how much to shift to get the ith uint
+        unpacked_data.append(((data >> shift_amt) & (nbits)).to(data.dtype))
+    return torch.stack(unpacked_data,dim=-1).view(*shape[:-1],shape[-1]*scale) # stack the unpacked data and reshape to the original shape
+
+def _pack_by_rows(data, container_size, data_size, device) -> torch.Tensor:
+
+    scale = container_size // data_size
+    assert scale > 1, f"container_size ({container_size}) is not larger than data_size ({data_size})"
+    assert data.shape[0] >= scale, f"not enough values to pack, data.shape[0] ({data.shape[0]}) < scale ({scale})"
+    # pad the data to be divisible by scale
+    if data.shape[0] % scale != 0:
+        padding = torch.zeros((scale - data.shape[0] % scale, *data.shape[1:],), dtype=data.dtype).to(device)
+        data = torch.cat([data, padding], dim=0).cuda()
+
+    shape = data.shape
+    ret = reduce(lambda x,y: x|y,[data[i::scale, ...] << container_size-data_size*(i+1) for i in range(scale)])
+    return ret.view(shape[0] // scale, *shape[1:]).to(device)
+
+def _pack_by_cols(data, container_size, data_size, device) -> torch.Tensor:
+    scale = container_size // data_size
+    assert scale > 1, f"container_size ({container_size}) not double the capacity ofdata_size ({data_size})"
+    # pad the data to be divisible by scale
+    if data.shape[-1] % scale != 0:
+        padding = torch.zeros((*data.shape[:-1], scale - data.shape[-1] % scale), dtype=data.dtype).to(device)
+        data = torch.cat([data, padding], dim=-1).cuda()
+
+    shape = data.shape
+    data = data.contiguous().view(-1)
+    #shift the data to the different indexes within the larger dtype and then union them together
+    ret = reduce(lambda x,y: x|y,[data[i::scale] << container_size-data_size*(i+1) for i in range(scale)])
+    return ret.view(*shape[:-1],shape[-1] // scale).to(device)