Skip to content

fix(tools): fp8_cast_bf16 upcast weights missing scale_inv to bf16 #1641

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

Open
nick-rui wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix(tools): fp8_cast_bf16 upcast weights missing scale_inv to bf16 #1641

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ca784e4
fix(tools): fp8_cast_bf16 upcast weights missing scale_inv to bf16
nick-rui Jul 12, 2026
88f7db8
fp8_cast_bf16: only treat 1-byte floating-point tensors as FP8
nick-rui Jul 12, 2026
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
4 changes: 2 additions & 2 deletions tools/fp8_cast_bf16.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,8 +100,8 @@ def get_tensor(tensor_name):
fp8_weight_names.append(weight_name)
new_state_dict[weight_name] = weight_dequant(weight, scale_inv)
except KeyError:
print(f"Warning: Missing scale_inv tensor for {weight_name}, skipping conversion")
new_state_dict[weight_name] = weight
print(f"Warning: Missing scale_inv tensor for {weight_name}, upcasting to {torch.get_default_dtype()}")
new_state_dict[weight_name] = weight.to(torch.get_default_dtype())
Comment on lines +103 to +104

@gemini-code-assist gemini-code-assist Bot Jul 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The check weight.element_size() == 1 on line 95 matches any 1-byte tensor, including torch.bool, torch.int8, and torch.uint8 (such as boolean masks or index tensors). If any such non-FP8 1-byte tensors do not have a corresponding _scale_inv tensor, they will trigger the KeyError and be incorrectly upcast to bfloat16. To prevent this, we should ensure that we only upcast actual floating-point (FP8) tensors by checking weight.is_floating_point().

Suggested change
print(f"Warning: Missing scale_inv tensor for {weight_name}, upcasting to {torch.get_default_dtype()}")
new_state_dict[weight_name] = weight.to(torch.get_default_dtype())
if weight.is_floating_point():
print(f"Warning: Missing scale_inv tensor for {weight_name}, upcasting to {torch.get_default_dtype()}")
new_state_dict[weight_name] = weight.to(torch.get_default_dtype())
else:
new_state_dict[weight_name] = weight

@nick-rui nick-rui Jul 12, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch. Fixed in 88f7db8 by guarding the FP8 branch condition itself (weight.element_size() == 1 and weight.is_floating_point()) instead of adding a check inside the except block — this keeps non-float 1-byte tensors on the pass-through path and also keeps them out of the dequant path entirely.

else:
new_state_dict[weight_name] = weight

Expand Down