Simulate group-wise quantization.
(tensor, config)
| 111 | |
| 112 | |
| 113 | def compress(tensor, config): |
| 114 | """Simulate group-wise quantization.""" |
| 115 | if not config.enabled: |
| 116 | return tensor |
| 117 | |
| 118 | group_size, num_bits, group_dim, symmetric = ( |
| 119 | config.group_size, |
| 120 | config.num_bits, |
| 121 | config.group_dim, |
| 122 | config.symmetric, |
| 123 | ) |
| 124 | assert num_bits <= 8 |
| 125 | |
| 126 | original_shape = tensor.shape |
| 127 | num_groups = (original_shape[group_dim] + group_size - 1) // group_size |
| 128 | new_shape = ( |
| 129 | original_shape[:group_dim] |
| 130 | + (num_groups, group_size) |
| 131 | + original_shape[group_dim + 1 :] |
| 132 | ) |
| 133 | |
| 134 | # Pad |
| 135 | pad_len = group_size - original_shape[group_dim] % group_size |
| 136 | if pad_len != 0: |
| 137 | pad_shape = ( |
| 138 | original_shape[:group_dim] + (pad_len,) + original_shape[group_dim + 1 :] |
| 139 | ) |
| 140 | tensor = torch.cat( |
| 141 | [tensor, torch.zeros(pad_shape, dtype=tensor.dtype, device=tensor.device)], |
| 142 | dim=group_dim, |
| 143 | ) |
| 144 | data = tensor.view(new_shape) |
| 145 | |
| 146 | # Quantize |
| 147 | if symmetric: |
| 148 | B = 2 ** (num_bits - 1) - 1 |
| 149 | scale = B / torch.max(data.abs(), dim=group_dim + 1, keepdim=True)[0] |
| 150 | data = data * scale |
| 151 | data = data.clamp_(-B, B).round_().to(torch.int8) |
| 152 | return data, scale, original_shape |
| 153 | else: |
| 154 | B = 2**num_bits - 1 |
| 155 | mn = torch.min(data, dim=group_dim + 1, keepdim=True)[0] |
| 156 | mx = torch.max(data, dim=group_dim + 1, keepdim=True)[0] |
| 157 | |
| 158 | scale = B / (mx - mn) |
| 159 | data = data - mn |
| 160 | data *= scale |
| 161 | |
| 162 | data = data.clamp_(0, B).round_().to(torch.uint8) |
| 163 | return data, mn, scale, original_shape |
| 164 | |
| 165 | |
| 166 | def decompress(packed_data, config): |
no outgoing calls
no test coverage detected