Fix disk offload crash on FP8 tensors

src/accelerate/utils/offload.py

@@ -29,6 +29,10 @@ def offload_weight(weight, weight_name, offload_folder, index=None):
         # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
         weight = weight.view(torch.int16)
         dtype = "bfloat16"
+    elif str(weight.dtype).startswith("torch.float8_"):
+        # NumPy does not handle any FP8 dtype either, so reinterpret the 1-byte data as int8.
+        dtype = str(weight.dtype).split(".")[1]
+        weight = weight.view(torch.int8)
     array = weight.cpu().numpy()
     tensor_file = os.path.join(offload_folder, f"{weight_name}.dat")
     if index is not None:
@@ -53,6 +57,9 @@ def load_offloaded_weight(weight_file, weight_info):
     if dtype == "bfloat16":
         # NumPy does not support bfloat16 so this was saved as a int16
         dtype = "int16"
+    elif dtype.startswith("float8_"):
+        # NumPy does not support any FP8 dtype either, so this was saved as an int8
+        dtype = "int8"
 
     weight = np.memmap(weight_file, dtype=dtype, shape=shape, mode="r")
 
@@ -61,6 +68,8 @@ def load_offloaded_weight(weight_file, weight_info):
     weight = torch.tensor(weight)
     if weight_info["dtype"] == "bfloat16":
         weight = weight.view(torch.bfloat16)
+    elif weight_info["dtype"].startswith("float8_"):
+        weight = weight.view(getattr(torch, weight_info["dtype"]))
 
     return weight
 

tests/test_offload.py

@@ -26,6 +26,7 @@
     offload_state_dict,
     offload_weight,
 )
+from accelerate.utils.versions import is_torch_version
 
 
 class ModelForTest(nn.Module):
@@ -56,16 +57,30 @@ def test_offload_state_dict(self):
     def test_offload_weight(self):
         dtypes = [torch.float16, torch.float32, torch.bfloat16]
 
+        # NumPy has no FP8 dtypes either, so these go through the same int8-view path as
+        # bfloat16 goes through int16. torch.randn(..., dtype=float8_*) isn't implemented, so
+        # the tensors are produced the way they actually show up in practice: cast down from a
+        # higher-precision tensor.
+        if is_torch_version(">=", "2.1.0"):
+            for name in ("float8_e4m3fn", "float8_e5m2"):
+                if hasattr(torch, name):
+                    dtypes.append(getattr(torch, name))
+
         for dtype in dtypes:
-            weight = torch.randn(2, 3, dtype=dtype)
+            is_fp8 = str(dtype).startswith("torch.float8_")
+            weight = torch.randn(2, 3, dtype=torch.float32).to(dtype) if is_fp8 else torch.randn(2, 3, dtype=dtype)
             with TemporaryDirectory() as tmp_dir:
                 index = offload_weight(weight, "weight", tmp_dir, {})
                 weight_file = os.path.join(tmp_dir, "weight.dat")
                 assert os.path.isfile(weight_file)
                 assert index == {"weight": {"shape": [2, 3], "dtype": str(dtype).split(".")[1]}}
 
                 new_weight = load_offloaded_weight(weight_file, index["weight"])
-                assert torch.equal(weight, new_weight)
+                assert new_weight.dtype == weight.dtype
+                # Compare on raw bits: FP8/bfloat16 round-trip through an int view, and float
+                # equality can't be trusted to catch a byte-level corruption anyway.
+                int_dtype = {1: torch.int8, 2: torch.int16, 4: torch.int32}[weight.element_size()]
+                assert torch.equal(weight.view(int_dtype), new_weight.view(int_dtype))
 
     def test_offload_weights_loader(self):
         model = ModelForTest()