Quantize a floating-point tensor to the target dtype with appropriate scaling. Args: tensor: Input tensor (float) dtype: Target dtype for quantization Returns: quantized_data: Quantized tensor scale: Scale factor used
(tensor, dtype)
| 63 | raise ValueError(f"Unsupported dtype for quantization: {dtype}") |
| 64 | |
| 65 | def quantize_tensor(tensor, dtype): |
| 66 | """ |
| 67 | Quantize a floating-point tensor to the target dtype with appropriate scaling. |
| 68 | |
| 69 | Args: |
| 70 | tensor: Input tensor (float) |
| 71 | dtype: Target dtype for quantization |
| 72 | Returns: |
| 73 | quantized_data: Quantized tensor |
| 74 | scale: Scale factor used |
| 75 | """ |
| 76 | scale = compute_scale_for_dtype(tensor, dtype) |
| 77 | |
| 78 | if dtype == torch.int8: |
| 79 | quantized_data = torch.clamp(torch.round(tensor / scale), -128, 127).to(dtype) |
| 80 | elif dtype == torch.uint8: |
| 81 | quantized_data = torch.clamp(torch.round(tensor / scale), 0, 255).to(dtype) |
| 82 | elif dtype in (torch.float8_e4m3fn, torch.float8_e5m2): |
| 83 | # For float8, we scale and then cast directly to the target type |
| 84 | # The casting operation will handle the appropriate rounding |
| 85 | scaled_tensor = tensor / scale |
| 86 | quantized_data = scaled_tensor.to(dtype) |
| 87 | else: |
| 88 | raise ValueError(f"Unsupported dtype for quantization: {dtype}") |
| 89 | |
| 90 | return quantized_data, scale |
| 91 | |
| 92 | |
| 93 | def update_parameter(target, result_float): |
no test coverage detected