Skip to content

Commit

Permalink
unbreak float8 static quant tutorial
Browse files Browse the repository at this point in the history
Summary:

The tutorial was broken, fixing it.

Test Plan:

```
python tutorials/calibration_flow/static_quant.py
```

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: aebfc44af855a8c39ddcf1eb06c6e13e4c12d19a
ghstack-comment-id: 2655560001
Pull Request resolved: #1709
  • Loading branch information
vkuzo committed Feb 14, 2025
1 parent 413689d commit 6c6299c
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 1 deletion.
1 change: 1 addition & 0 deletions torchao/dtypes/floatx/float8_layout.py
Original file line number Diff line number Diff line change
Expand Up @@ -253,6 +253,7 @@ def _linear_fp8_act_fp8_weight_impl(
):
"""Implements matmul between FP8 input and FP8 weight with compute using _scaled_mm"""
scaled_mm_config = weight_tensor._layout.mm_config
assert scaled_mm_config is not None
out_shape = get_out_shape(input_tensor.shape, weight_tensor.shape)

# Weight tensor preprocessing
Expand Down
3 changes: 2 additions & 1 deletion tutorials/calibration_flow/static_quant.py
Original file line number Diff line number Diff line change
Expand Up @@ -163,12 +163,13 @@ def __init__(
weight, weight_scale, weight_zero_point, block_size, self.target_dtype
)
elif self.target_dtype == torch.float8_e4m3fn:
mm_config = Float8MMConfig(use_fast_accum=True)
self.qweight = to_affine_quantized_floatx_static(
weight,
weight_scale,
block_size,
target_dtype,
Float8Layout(mm_config=None),
Float8Layout(mm_config=mm_config),
)
else:
raise ValueError(f"Unsupported target dtype {self.target_dtype}")
Expand Down

0 comments on commit 6c6299c

Please sign in to comment.