Update dtype_byte_size to handle torch.float8_e4m3fn/float8_e5m2 types

[mgoin]

What does this PR do?

Currently using the new torch.float8_e4m3fn dtype will cause an error because dtype_byte_size() doesn't know where to find the number of bits in the dtype string.

Code to reproduce:

from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained("echarlaix/tiny-random-mistral")
model.lm_head.weight = torch.nn.Parameter(model.lm_head.weight.to(torch.float8_e4m3fn))
model.save_pretrained("test")

Error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/michael/venvs/test/lib/python3.10/site-packages/transformers/modeling_utils.py", line 2557, in save_pretrained
    shards, index = shard_checkpoint(state_dict, max_shard_size=max_shard_size, weights_name=weights_name)
  File "/home/michael/venvs/test/lib/python3.10/site-packages/transformers/modeling_utils.py", line 381, in shard_checkpoint
    weight_size = weight.numel() * dtype_byte_size(weight.dtype)
  File "/home/michael/venvs/test/lib/python3.10/site-packages/transformers/modeling_utils.py", line 328, in dtype_byte_size
    raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
ValueError: `dtype` is not a valid dtype: torch.float8_e4m3fn.

We can fix this by changing the regex from [^\d](\d+)$ to [^\d](\d+)_? so we match on every set of numbers following an alphabetical character. This can match multiple groups, but we can always look at the first group to get the first number which is significant. See the picture below for a demonstration of the matches on all numbered pytorch dtypes:

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Tagging @sgugger since they touched this function last.

[robertgshaw2-redhat]

Thanks! This is great for fp8

[amyeroberts]

Thanks for fixing this!

Could you add a test?

[mgoin]

Sure @amyeroberts, I added a test specifically for dtype_byte_size over most of the torch dtypes. Let me know if you have something else in mind, thanks!

[amyeroberts]

Beautiful - thanks for updating and adding tests!

[amyeroberts]

For the quality checks, running make fixup should resolve