Add FP8 Adam (#482)

* update benchmark

* add rank1 option to lpmm

* add comma

* update readme

* remove unwanted file

* update

* add Adam fp8

* add FP8 AdamW and test

* update readme

* change reason to xfail, since 2.2 also have float8

* at guard for FP8

* update readme

* fix guard

benchmarks/benchmark_low_bit_adam.py

@@ -28,15 +28,16 @@
 from torchvision.transforms import v2
 from tqdm import tqdm
 
-from torchao.prototype.low_bit_optim import Adam4bit, Adam8bit
+from torchao.prototype import low_bit_optim
 
 # lpmm doesn't have Adam, only AdamW
 OPTIM_MAP = dict(
  Adam=torch.optim.Adam,
  Adam8bitBnb=bnb.optim.Adam8bit,
- Adam8bitAo=Adam8bit,
+ Adam8bitAo=low_bit_optim.Adam8bit,
+ AdamFp8Ao=low_bit_optim.AdamFp8,
  Adam4bitLpmm=partial(lpmm.optim.AdamW, weight_decay=0, fused=True),
- Adam4bitAo=Adam4bit,
+ Adam4bitAo=low_bit_optim.Adam4bit,
  Adam4bitRank1Lpmm=partial(lpmm.optim.AdamW, weight_decay=0, qconfig=argparse.Namespace(scale_type="rank1")),
 )
 

test/prototype/test_low_bit_optim.py

@@ -139,6 +139,22 @@ def test_optim_4bit_correctness(self, optim_name):
  for p1, p2 in zip(model1.parameters(), model2.parameters()):
  torch.testing.assert_close(p2, p1, rtol=1e-5, atol=1e-5)
 
+ @pytest.mark.xfail(not TORCH_VERSION_AFTER_2_3, reason="torch.compile() fails for PyTorch < 2.3")
+ @parametrize("optim_name", ["AdamFp8", "AdamWFp8"])
+ @parametrize("device", _DEVICES)
+ def test_optim_fp8_smoke(self, optim_name, device):
+ if device == "cuda" and torch.cuda.get_device_capability() < (8, 9):
+ pytest.skip("FP8 requires compute capability >= 8.9")
+
+ model = nn.Sequential(nn.Linear(32, 1024), nn.ReLU(), nn.Linear(1024, 128)).to(device)
+ optim = getattr(low_bit_optim, optim_name)(model.parameters())
+
+ x = torch.randn(4, 32, device=device)
+ loss = model(x).sum()
+ loss.backward()
+ optim.step()
+ optim.zero_grad()
+
 
 instantiate_parametrized_tests(TestQuantize)
 instantiate_parametrized_tests(TestOptim)

torchao/prototype/low_bit_optim/README.md

@@ -4,6 +4,7 @@ This folder implements:
 
 - 8-bit optimizers as outlined in https://arxiv.org/abs/2110.02861
 - 4-bit optimizers as outlined in https://arxiv.org/abs/2309.01507
+- FP8 optimizers using the native `torch.float8_e4m3fn` dtype (experimental)
 
 The implementation is fully done in Python (with tensor subclass) and relies on `torch.compile()` to generate efficient fused kernel.
 
@@ -18,12 +19,12 @@ model = ...
 optim = Adam8bit(model.parameters())
 ```
 
-To use 4-bit Adam, replace the above with `Adam4bit`. You can also change quantization block size by passing `block_size=value` to the optimizer. By default, block size is 2048 for 8-bit optimizers, and 128 for 4-bit optimizers.
+To use 4-bit Adam, replace the above with `Adam4bit`. Similarly for `AdamFp8`. You can also change quantization block size by passing `block_size=value` to the optimizer. By default, block size is 2048 for 8-bit and FP8 optimizers, and 128 for 4-bit optimizers.
 
-**Other optimizers**: AdamW is also available as `AdamW8bit` and `AdamW4bit`. Other optimizers can be added based on demand.
+**Other optimizers**: AdamW is also available as `AdamW8bit`, `AdamW4bit`, and `AdamWFp8`. Other optimizers can be added based on demand.
 
 NOTE:
-- The low-bit optimizers require PyTorch >= 2.3
+- The low-bit optimizers require PyTorch >= 2.3. FP8 optimizers require CUDA compute capability >= 8.9.
 - For 4-bit optimizers, we don't implement rank-1 normalization for quantizing 2nd moment as originally done in the paper.
 - **Known issue**: When learning rate is updated every step (e.g. using cosine learning rate scheduler), training speed is slower. This is because we have to convert learning rate to a CUDA tensor (which incurs expensive memory transfer cost), since torch.compile() will treat a Python float as a constant and trigger recompile whenever the value is changed.
 
@@ -38,6 +39,7 @@ Adam impl | max memory (GB) | time taken for 2nd epoch | accuracy
 PyTorch | 12.94 | 8m 18s | 91.14
 bnb 8-bit | 8.31 | 6m 50s | 90.67
 ao 8-bit | 8.32 | 9m 04s | 90.71
+ao FP8 E4M3 | 8.32 | 6m 38s | 91.08
 lpmm 4-bit | 7.72 | 5m 59s | 89.97
 ao 4-bit | 7.72 | 7m 00s | 89.94
 lpmm 4-bit (*) | 7.73 | 11m 10s | 89.71

torchao/prototype/low_bit_optim/__init__.py

@@ -1,2 +1,2 @@
-from .adam import Adam8bit, Adam4bit
-from .adamw import AdamW8bit, AdamW4bit
+from .adam import Adam8bit, Adam4bit, AdamFp8
+from .adamw import AdamW8bit, AdamW4bit, AdamWFp8

torchao/prototype/low_bit_optim/adam.py

@@ -6,6 +6,7 @@
 
 from .subclass_8bit import maybe_new_8bit_zero_buffer
 from .subclass_4bit import maybe_new_4bit_zero_buffer
+from .subclass_fp8 import maybe_new_fp8_zero_buffer
 
 
 class _Adam(Optimizer):
@@ -155,3 +156,22 @@ def __init__(
  super().__init__(params, lr, betas, eps, weight_decay, amsgrad, block_size=block_size)
 
  _new_buffer = staticmethod(maybe_new_4bit_zero_buffer)
+
+
+class AdamFp8(_Adam):
+ def __init__(
+ self,
+ params,
+ lr=1e-3,
+ betas=(0.9, 0.999),
+ eps=1e-8,
+ weight_decay=0,
+ amsgrad=False,
+ *,
+ block_size=2048
+ ) -> None:
+ super().__init__(params, lr, betas, eps, weight_decay, amsgrad, block_size=block_size)
+
+ @staticmethod
+ def _new_buffer(p: Tensor, signed: bool, block_size: int):
+ return maybe_new_fp8_zero_buffer(p, block_size)

torchao/prototype/low_bit_optim/adamw.py

@@ -6,6 +6,7 @@
 
 from .subclass_8bit import maybe_new_8bit_zero_buffer
 from .subclass_4bit import maybe_new_4bit_zero_buffer
+from .subclass_fp8 import maybe_new_fp8_zero_buffer
 
 
 class _AdamW(Optimizer):
@@ -154,3 +155,22 @@ def __init__(
  super().__init__(params, lr, betas, eps, weight_decay, amsgrad, block_size=block_size)
 
  _new_buffer = staticmethod(maybe_new_4bit_zero_buffer)
+
+
+class AdamWFp8(_AdamW):
+ def __init__(
+ self,
+ params,
+ lr=1e-3,
+ betas=(0.9, 0.999),
+ eps=1e-8,
+ weight_decay=1e-2,
+ amsgrad=False,
+ *,
+ block_size=2048
+ ) -> None:
+ super().__init__(params, lr, betas, eps, weight_decay, amsgrad, block_size=block_size)
+
+ @staticmethod
+ def _new_buffer(p: Tensor, signed: bool, block_size: int):
+ return maybe_new_fp8_zero_buffer(p, block_size)

torchao/prototype/low_bit_optim/subclass_fp8.py

@@ -0,0 +1,106 @@
+import torch
+from torch import Tensor
+from torchao.dtypes.utils import _implements, _ATEN_OP_OR_TORCH_FN_TABLE
+
+
+aten = torch.ops.aten
+DTYPE = torch.float8_e4m3fn
+
+
+def quantize_fp8(input: Tensor, block_size: int):
+ input = input.view(-1, block_size)
+ 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, scale
+
+
+class OptimStateFp8(Tensor):
+ implements = classmethod(_implements)
+ tensor_attrs = ["codes", "scale"]
+
+ @staticmethod
+ def __new__(cls, codes: Tensor, scale: Tensor):
+ return Tensor._make_wrapper_subclass(
+ cls,
+ codes.shape,
+ device=codes.device,
+ requires_grad=False,
+ )
+
+ def __init__(self, codes: Tensor, scale: Tensor):
+ assert codes.dtype is DTYPE
+ self.codes = codes
+ self.scale = scale
+
+ @property
+ def block_size(self):
+ return self.codes.numel() // self.scale.numel()
+
+ def __tensor_flatten__(self):
+ return self.tensor_attrs, []
+
+ @classmethod
+ def __tensor_unflatten__(cls, tensor_data_dict, tensor_attributes, outer_size=None, outer_stride=None):
+ return cls(*[tensor_data_dict[name] for name in cls.tensor_attrs], *tensor_attributes)
+
+ 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)
+
+ dtype = output_dtype or torch.get_default_dtype()
+ return float_data.view(self.codes.shape).to(dtype)
+
+ @classmethod
+ def zeros(cls, shape, block_size: int = 2048, device=None):
+ codes = torch.zeros(shape, dtype=DTYPE, device=device)
+ scale = torch.zeros(codes.numel() // block_size, device=device)
+ return cls(codes, scale)
+
+ def __repr__(self):
+ return (
+ f"{self.__class__.__name__}(block_size={self.block_size}, "
+ f"shape={tuple(self.shape)}, device={self.device}, requires_grad={self.requires_grad})"
+ )
+
+ @classmethod
+ def __torch_dispatch__(cls, func, types, args, kwargs):
+ if func in _ATEN_OP_OR_TORCH_FN_TABLE[cls]:
+ return _ATEN_OP_OR_TORCH_FN_TABLE[cls][func](func, *args, **kwargs)
+
+ raise NotImplementedError(f"{cls.__name__} dispatch: attempting to run {func}, this is not supported")
+
+
+@OptimStateFp8.implements(aten.copy_.default)
+def _(func, *args, **kwargs):
+ dst = args[0]
+ src = args[1]
+
+ if isinstance(dst, OptimStateFp8) and isinstance(src, OptimStateFp8):
+ assert dst.block_size == src.block_size
+ dst.codes.copy_(src.codes)
+ dst.scale.copy_(src.scale)
+
+ elif isinstance(dst, OptimStateFp8):
+ codes, scale = quantize_fp8(src, dst.block_size)
+ dst.codes.copy_(codes)
+ dst.scale.copy_(scale)
+
+ else:
+ dst.copy_(src.dequantize())
+
+ return dst
+
+
+@OptimStateFp8.implements(aten.lerp.Scalar)
+def _(func, *args, **kwargs):
+ args = [x.dequantize() if isinstance(x, OptimStateFp8) else x for x in args]
+ return func(*args, **kwargs)
+
+
+def maybe_new_fp8_zero_buffer(p: Tensor, block_size: int = 2048):
+ if p.numel() >= 4096 and p.numel() % block_size == 0:
+ out = OptimStateFp8.zeros(p.shape, block_size, device=p.device)
+ else:
+ out = torch.zeros_like(p)
+ return out