Fix dtype_byte_size reporting zero bytes for sub-byte and packed dtypes

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}.")

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)