Add fp8_e5m2 support in dtype_byte_size (#3625)

* float8_e5m2 device_map

* remove prints

Author: SunMarc

src/accelerate/utils/modeling.py

@@ -169,7 +169,7 @@ 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 == torch.float8_e4m3fn:
+    elif is_torch_version(">=", "2.1.0") and dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
         return 1
     bit_search = re.search(r"[^\d](\d+)$", str(dtype))
     if bit_search is None: