Skip to content
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

[float8] Re-enable slow-accum in the bwd of axis-wise scaling schemes #1325

Merged
merged 3 commits into from
Dec 4, 2024
Merged

[float8] Re-enable slow-accum in the bwd of axis-wise scaling schemes #1325

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c9e26bd
Update
lw Nov 22, 2024
7debcd9
Update
lw Nov 22, 2024
6f4615b
Update
lw Dec 4, 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
23 changes: 0 additions & 23 deletions torchao/float8/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,6 @@ class Float8LinearConfig:
#
# Per-gemm configuration for gemms calculating `output`, `grad_input` and
# `grad_weight`
# TODO(this PR): throw warning if fast_accum False is used with axiswise scaling
#
gemm_config_output: Float8GemmConfig = Float8GemmConfig(use_fast_accum=True)
gemm_config_grad_input: Float8GemmConfig = Float8GemmConfig()
Expand Down Expand Up @@ -317,21 +316,10 @@ def recipe_name_to_linear_config(
cc_w = CastConfig(scaling_granularity=ScalingGranularity.AXISWISE)
cc_go = CastConfig(scaling_granularity=ScalingGranularity.AXISWISE)

# The current rowwise CUTLASS kernels in `torch._scaled_mm` are only
# fast with `use_fast_accum=True`. Note that rowwise scaling is more
# accurate than tensorwise scaling, so the overall impact on accuracy
# of tensorwise vs rowwise taking this flag into account will vary.
gc_o = Float8GemmConfig(use_fast_accum=True)
gc_gi = Float8GemmConfig(use_fast_accum=True)
gc_gw = Float8GemmConfig(use_fast_accum=True)

return Float8LinearConfig(
cast_config_input=cc_i,
cast_config_weight=cc_w,
cast_config_grad_output=cc_go,
gemm_config_output=gc_o,
gemm_config_grad_input=gc_gi,
gemm_config_grad_weight=gc_gw,
)

elif recipe_name is Float8LinearRecipeName.LW_AXISWISE_WITH_GW_HP:
Expand Down Expand Up @@ -359,24 +347,13 @@ def recipe_name_to_linear_config(
cc_i_gw = CastConfig(scaling_type=ScalingType.DISABLED)
cc_go_gw = CastConfig(scaling_type=ScalingType.DISABLED)

# The current rowwise CUTLASS kernels in `torch._scaled_mm` are only
# fast with `use_fast_accum=True`. Note that rowwise scaling is more
# accurate than tensorwise scaling, so the overall impact on accuracy
# of tensorwise vs rowwise taking this flag into account will vary.
gc_o = Float8GemmConfig(use_fast_accum=True)
gc_gi = Float8GemmConfig(use_fast_accum=True)
gc_gw = Float8GemmConfig(use_fast_accum=True)

return Float8LinearConfig(
cast_config_input=cc_i,
cast_config_weight=cc_w,
cast_config_grad_output=cc_go,
cast_config_input_for_grad_weight=cc_i_gw,
cast_config_weight_for_grad_input=cc_w_gi,
cast_config_grad_output_for_grad_weight=cc_go_gw,
gemm_config_output=gc_o,
gemm_config_grad_input=gc_gi,
gemm_config_grad_weight=gc_gw,
)

else:
Expand Down
36 changes: 24 additions & 12 deletions torchao/float8/float8_python_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,19 +37,25 @@ def addmm_float8_unwrapped(
a_inverse_scale = a_scale.reciprocal()
b_inverse_scale = b_scale.reciprocal()

if output_dtype == torch.float32 and bias is not None:
post_inverse_scale = None
if (
a_scale.shape == (a_data.shape[0], 1)
and b_scale.shape == (1, b_data.shape[1])
and not use_fast_accum
):
# The rowwise CUTLASS-based kernel is so slow without fast-accum that
Copy link
Contributor

Choose a reason for hiding this comment

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

just curious, do we have any OSS shareable evidence (perf/accuracy) on doing this versus rowwise with fast-accum off that we can add here?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Screenshot 2024-12-03 at 19 09 19

I ran a quick benchmark on my H100 with a recent-ish version of PyTorch (nightly from Nov 12). I samples all MxNxK matmul shapes where each of M, N and K is a power of two between 512 and 16384. Here I'm plotting the slowdowns observed when activating slow-accum for the rowwise (CUTLASS-based) and tensorwise (cuBLAS-based) modes

In summary: in tensorwise we get a max slowdown of 50% (usually much less), with rowwise we typically are 2x as slow, with peaks of 4.5x as slow as fast-accum.

(I suspect that for very small shapes the benchmark was CPU-bound hence slow-accum looks as fast as fast-accum, but that's probably misleading)

# we'd rather use the tensorwise cuBLAS-based kernel and do the scaling
# manually afterwards (hoping Inductor will be able to fuse it).
post_inverse_scale = a_inverse_scale * b_inverse_scale
a_inverse_scale = a_inverse_scale.new_ones(())
b_inverse_scale = a_inverse_scale.new_ones(())

post_bias = None
if output_dtype == torch.float32:
# Bias is not supported by _scaled_mm when output is fp32
output = torch._scaled_mm(
a_data,
b_data,
scale_a=a_inverse_scale,
scale_b=b_inverse_scale,
scale_result=output_scale,
out_dtype=output_dtype,
use_fast_accum=use_fast_accum,
)
output += bias
return output
post_bias = bias
bias = None

output = torch._scaled_mm(
a_data,
b_data,
Expand All @@ -60,4 +66,10 @@ def addmm_float8_unwrapped(
out_dtype=output_dtype,
use_fast_accum=use_fast_accum,
)

if post_inverse_scale is not None:
output *= post_inverse_scale
if post_bias is not None:
output += post_bias

return output
Loading