| 56 | self.scale = None |
| 57 | |
| 58 | def quantize(self, |
| 59 | input, |
| 60 | q_bits=8, |
| 61 | q_mantisa_bits=3, |
| 62 | stochastic_mode=False, |
| 63 | return_meta_tensor=False) -> torch.Tensor: |
| 64 | assert input.dtype == torch.bfloat16, "only support bf16 for now" |
| 65 | if return_meta_tensor: |
| 66 | assert q_bits == 8, "meta tensor is only supported with q_bit=8" |
| 67 | |
| 68 | self.orig_dtype = input.dtype |
| 69 | self.orig_shape = input.shape |
| 70 | |
| 71 | if q_bits == 8: |
| 72 | pass |
| 73 | elif q_bits == 12: |
| 74 | q_mantisa_bits = 4 |
| 75 | elif q_bits == 6: |
| 76 | q_mantisa_bits = 2 |
| 77 | elif q_bits == 4: |
| 78 | q_mantisa_bits = 1 |
| 79 | else: |
| 80 | assert (0), \ |
| 81 | f"Missing {q_bits}-quantization, please add the template arguments for the kernel to support this precision!" |
| 82 | self.num_groups = input.numel() // self.group_size |
| 83 | self.input_q = torch.ones(self.num_groups, |
| 84 | int(self.group_size * q_bits) // 8 + 4, |
| 85 | dtype=torch.uint8, |
| 86 | device=input.device) |
| 87 | out = fp_quant_module.quantize(self.input_q, input, self.group_size, stochastic_mode, q_bits, q_mantisa_bits) |
| 88 | if return_meta_tensor: |
| 89 | data, self.scale = out.split(self.group_size, dim=-1) |
| 90 | data = data.contiguous().reshape(input.shape) |
| 91 | self.scale = self.scale.contiguous() |
| 92 | del self.input_q |
| 93 | del out |
| 94 | gc.collect() |
| 95 | get_accelerator().empty_cache() |
| 96 | return data, self.scale |
| 97 | |
| 98 | return out |
| 99 | |
| 100 | def to(self, *args, **kwargs): |
| 101 | # Intermediate tensors may need to be moved to different devices |