[low-bit optim] Add coat for float8 optimizer

torchao/prototype/low_bit_optim/subclass_fp8.py

@@ -1,5 +1,6 @@
 import torch
 from torch import Tensor
+from typing import Optional
 from torch.utils._python_dispatch import return_and_correct_aliasing
 
 from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, TorchAOBaseTensor
@@ -11,13 +12,26 @@
 DTYPE = torch.float8_e4m3fn
 
 
-def quantize_fp8(input: Tensor, block_size: int):
+def quantize_fp8(input: Tensor, block_size: int, dynamic_range_expansion: bool):
+
     shape = input.shape
     input = input.view(-1, block_size)
+    k =  aten.ones(input.shape[0], dtype=DTYPE, device=input.device) if dynamic_range_expansion else None
+
+    if dynamic_range_expansion:
+        # NOTE: the calculation is from the paper https://arxiv.org/abs/2410.19313
+        # The idea is to align optimizer state distributions more closely 
+        # with the FP8 representation range, reducing the quantization error.
+        Rdtype = torch.finfo(DTYPE).max / torch.finfo(DTYPE).min # calculate the range of the dtype
+        Rx = input.abs().amax(-1).clip(1e-12) / input.abs().amin(-1).clip(1e-12) # range of input max and min
+        k = torch.log(Rdtype) / torch.log(Rx)# calculating optimal value k dynamically
+        input = input.sign() * (input.abs() ** k.view(-1, 1))
+
     scale = input.abs().amax(-1).clip(1e-12) / torch.finfo(DTYPE).max
     input = input / scale.view(-1, 1)
     codes = input.to(DTYPE).view(-1)
-    return codes.view(shape), scale
+
+    return codes.view(shape), scale, k
 
 
 # NOTE: FP8 sign bit is redundant for unsigned optim state.
@@ -30,7 +44,7 @@ class OptimStateFp8(TorchAOBaseTensor):
     def __new__(cls, codes: Tensor, scale: Tensor):
         return Tensor._make_wrapper_subclass(cls, codes.shape, device=codes.device)
 
-    def __init__(self, codes: Tensor, scale: Tensor):
+    def __init__(self, codes: Tensor, scale: Tensor, k: Optional[Tensor] =None):
         """Create quantized FP8 optimizer state.
 
         Args
@@ -45,6 +59,7 @@ def __init__(self, codes: Tensor, scale: Tensor):
         assert scale.ndim == 1
         self.codes = codes
         self.scale = scale
+        self.k = k
         self.block_size = codes.numel() // scale.numel()
 
     def __tensor_flatten__(self):
@@ -62,15 +77,22 @@ def dequantize(self, output_dtype=None):
         float_data = self.codes.float()
         float_data = float_data.view(-1, self.block_size) * self.scale.view(-1, 1)
 
+        if self.k:
+            float_data = float_data.view(-1, self.block_size)
+            float_data =  float_data ** (1 / self.k.view(-1, 1))
+
         if output_dtype is not None:
             float_data = float_data.to(output_dtype)
+
         return float_data.view(self.codes.shape)
 
     @classmethod
-    def zeros(cls, shape, block_size: int = 256, device=None):
+    def zeros(cls, shape, block_size: int = 256, device=None, dynamic_range_expansion: bool = False):
+
         codes = torch.zeros(shape, dtype=DTYPE, device=device)
         scale = torch.zeros(codes.numel() // block_size, device=device)
-        return cls(codes, scale)
+        k = torch.ones(codes.numel() // block_size, device=device) if dynamic_range_expansion else None
+        return cls(codes, scale, k)
 
     def __repr__(self):
         return (
@@ -88,12 +110,15 @@ def _(func, types, args, kwargs):
         assert dst.block_size == src.block_size
         dst.codes.copy_(src.codes)
         dst.scale.copy_(src.scale)
+        dst.k.copy_(src.k)
 
     elif isinstance(dst, OptimStateFp8):
-        codes, scale = quantize_fp8(src, dst.block_size)
+
+        codes, scale, k = quantize_fp8(src, dst.block_size, True if dst.k is not None else False)
+
         dst.codes.copy_(codes)
         dst.scale.copy_(scale)
-
+        dst.k.copy_(k)
     else:
         dst.copy_(src.dequantize())
 
@@ -148,4 +173,4 @@ def _(func, types, args, kwargs):
 if TORCH_VERSION_AT_LEAST_2_5:
     from torch.serialization import add_safe_globals
 
-    add_safe_globals([OptimStateFp8])
+    add_safe_globals([OptimStateFp8])