From 82816c49b79fce2fdb629eaf2aa482a022e23c57 Mon Sep 17 00:00:00 2001 From: VihaanAgarwal Date: Wed, 9 Sep 2026 09:17:06 -0400 Subject: [PATCH] Fix dtype_byte_size reporting zero bytes for sub-byte and packed dtypes --- src/accelerate/utils/modeling.py | 15 +++------------ tests/test_modeling_utils.py | 4 ++++ 2 files changed, 7 insertions(+), 12 deletions(-) diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py index 6479fc242a8..f4c0e2d2fec 100644 --- a/src/accelerate/utils/modeling.py +++ b/src/accelerate/utils/modeling.py @@ -169,18 +169,9 @@ def dtype_byte_size(dtype: torch.dtype): return 1 / 2 elif dtype == CustomDtype.FP8: return 1 - elif is_torch_version(">=", "2.1.0") and dtype in [ - getattr(torch, name) - for name in ( - "float8_e4m3fn", - "float8_e5m2", - "float8_e4m3fnuz", - "float8_e5m2fnuz", - "float8_e8m0fnu", - ) - if hasattr(torch, name) - ]: - return 1 + elif is_torch_version(">=", "2.1.0"): + # The name regex below misreads FP8 and sub-byte dtypes such as `uint4` and `float4_e2m1fn_x2` + return dtype.itemsize bit_search = re.search(r"[^\d](\d+)$", str(dtype)) if bit_search is None: raise ValueError(f"`dtype` is not a valid dtype: {dtype}.") diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py index a7aa6e4c590..a34b8960946 100644 --- a/tests/test_modeling_utils.py +++ b/tests/test_modeling_utils.py @@ -153,6 +153,10 @@ def test_dtype_byte_size(self): ): if hasattr(torch, name): self.assertEqual(dtype_byte_size(getattr(torch, name)), 1, msg=name) + # Sub-byte and packed dtypes still occupy one byte per element in storage. + for name in ("uint4", "float4_e2m1fn_x2"): + if hasattr(torch, name): + self.assertEqual(dtype_byte_size(getattr(torch, name)), 1, msg=name) def check_set_module_tensor_for_device(self, model, device1, device2): assert model.linear1.weight.device == torch.device(device1)