diff --git a/src/accelerate/utils/offload.py b/src/accelerate/utils/offload.py index da8cfaf3ebf..48c26fe5f7b 100644 --- a/src/accelerate/utils/offload.py +++ b/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 diff --git a/tests/test_offload.py b/tests/test_offload.py index 95566a4bcbf..ffe244eb95f 100644 --- a/tests/test_offload.py +++ b/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,8 +57,18 @@ 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") @@ -65,7 +76,11 @@ def test_offload_weight(self): 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()