Skip to content

Fix float8_e4m3fn in modeling_utils #32193

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ArthurZucker merged 5 commits into from
Jul 24, 2024
+5 −2
Merged

Fix float8_e4m3fn in modeling_utils #32193

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7daf64c
Fix float8_e4m3fn in modeling_utils
SunMarc Jul 24, 2024
2686c0b
style
SunMarc Jul 24, 2024
9188822
Merge remote-tracking branch 'upstream/main' into fix-float8-dtype
SunMarc Jul 24, 2024
f70f000
fix
SunMarc Jul 24, 2024
ca67fba
comment
SunMarc Jul 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions src/transformers/modeling_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -855,6 +855,8 @@ def _load_state_dict_into_meta_model(
for old_key, new_key in zip(old_keys, new_keys):
state_dict[new_key] = state_dict.pop(old_key)

is_torch_e4m3fn_available = hasattr(torch, "float8_e4m3fn")

for param_name, param in state_dict.items():
# First part of the test is always true as load_state_dict_keys always contains state_dict keys.
if param_name not in loaded_state_dict_keys or param_name not in expected_keys:
Expand All @@ -866,9 +868,10 @@ def _load_state_dict_into_meta_model(
module_name = param_name
set_module_kwargs = {}

# We convert floating dtypes to the `dtype` passed. We want to keep the buffers/params
# We convert floating dtypes to the `dtype` passed except for float8_e4m3fn type. We also want to keep the buffers/params
# in int/uint/bool and not cast them.
if dtype is not None and torch.is_floating_point(param) and param.dtype != torch.float8_e4m3fn:
is_param_float8_e4m3fn = is_torch_e4m3fn_available and param.dtype == torch.float8_e4m3fn
if dtype is not None and torch.is_floating_point(param) and not is_param_float8_e4m3fn:
if (
keep_in_fp32_modules is not None
and any(
Expand Down