Spaces:
Running on Zero
Running on Zero
| """ComfyUI-style float8 weight storage as an Ostris backend. | |
| Matches the comfy_quant ``{"format": "float8_e4m3fn", | |
| "full_precision_matrix_mult": true}`` layout: the weight stored as | |
| torch.float8_e4m3fn plus one fp32 per-tensor scale, matmuls running on the | |
| dequantized weight (W8A16 numerics). Used both to import comfy fp8/fp8-mixed | |
| checkpoints and to quantize/export in that format. | |
| """ | |
| from typing import Optional | |
| import torch | |
| from toolkit.util.ostris_quant import OstrisLinear, OstrisQuantizer | |
| FLOAT8_QTYPES = ["float8_e4m3fn"] | |
| F8_MAX = torch.finfo(torch.float8_e4m3fn).max | |
| class Float8Quantizer(OstrisQuantizer): | |
| """fp8_e4m3 weight + fp32 per-tensor scale, dequantized matmul.""" | |
| def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None: | |
| scale = (weight_fp32.abs().max() / F8_MAX).clamp(min=1e-12) | |
| q = (weight_fp32 / scale).clamp(-F8_MAX, F8_MAX).to(torch.float8_e4m3fn) | |
| self.attach_(module, q, scale) | |
| def attach_( | |
| module: torch.nn.Module, | |
| qweight: torch.Tensor, # float8_e4m3fn (out, in) | |
| scale: torch.Tensor, # fp32 scalar | |
| ) -> None: | |
| """Register the quantized representation on the module. Used both by | |
| quantize_ and by importers of pre-quantized checkpoints.""" | |
| module.register_buffer("f8_qdata", qweight.contiguous(), persistent=False) | |
| module.register_buffer( | |
| "f8_scale", | |
| scale.detach().float().clone().reshape(1).view(torch.uint8), | |
| persistent=False, | |
| ) | |
| def dequantize(self, module: "OstrisLinear") -> torch.Tensor: | |
| scale = module.f8_scale.view(torch.float32)[0] | |
| return module.f8_qdata.to(torch.float32) * scale | |
| def requantize_(self, module: "OstrisLinear", fp_weight: torch.Tensor) -> None: | |
| w = fp_weight.to(torch.float32) | |
| scale = (w.abs().max() / F8_MAX).clamp(min=1e-12) | |
| module.f8_qdata.copy_((w / scale).clamp(-F8_MAX, F8_MAX).to(torch.float8_e4m3fn)) | |
| module.f8_scale.copy_(scale.reshape(1).view(torch.uint8)) | |