Some changes in `inner-padding` option #858

y-sq · 2024-09-09T22:30:20Z

Summary:
The diff modifies the padding option and added tests with compile:

For the scaled_mm of shape MxKxN, the current inner_padding option only pads the K dimension. However, if N is not divisible by 16, we also got the error

E       RuntimeError: CUDA error: CUBLAS_STATUS_NOT_SUPPORTED when calling `cublasLtMatmulAlgoGetHeuristic( ltHandle, computeDesc.descriptor(), Adesc.descriptor(), Bdesc.descriptor(), Cdesc.descriptor(), Ddesc.descriptor(), preference.descriptor(), 1, &heuristicResult, &returnedResult)`

So, modified the pad_inner option to also pad the N dimensions.

The compile of inner-padding only works with the triton PR Support FP8 constant triton-lang/triton#4222.

Before the triton PR, the inductor code-gen kernel fails at

tmp10 = tl.where(tmp6, tmp8, tmp9)

TypeError: unexpected type fp8e5 and fp8e5

Reviewed By: irobert0126

Differential Revision: D62003827

pytorch-bot · 2024-09-09T22:30:23Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/858

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 9b6212f with merge base 8aa6533 ():

NEW FAILURE - The following job has failed:

Run Float8 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
RuntimeError: Command docker exec -t e50b322cfcf1cd127d433e98a31e5aaa2b7eafbb18d4fc75996fda16486b8df4 /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

facebook-github-bot · 2024-09-09T22:30:42Z

This pull request was exported from Phabricator. Differential Revision: D62003827

facebook-github-bot · 2024-09-10T21:06:50Z

This pull request was exported from Phabricator. Differential Revision: D62003827

Summary: Pull Request resolved: pytorch#858 Add test cases to verify that the compile of inner-padding works with the triton PR triton-lang/triton#4222. Before the triton PR, the inductor code-gen kernel fails at ``` tmp10 = tl.where(tmp6, tmp8, tmp9) TypeError: unexpected type fp8e5 and fp8e5 ``` Reviewed By: irobert0126 Differential Revision: D62003827

facebook-github-bot · 2024-09-10T21:12:26Z

This pull request was exported from Phabricator. Differential Revision: D62003827

Summary: Pull Request resolved: pytorch#858 Add test cases to verify that the compile of inner-padding works with the triton PR triton-lang/triton#4222. Before the triton PR, the inductor code-gen kernel fails at ``` tmp10 = tl.where(tmp6, tmp8, tmp9) TypeError: unexpected type fp8e5 and fp8e5 ``` Reviewed By: irobert0126 Differential Revision: D62003827

Summary: Pull Request resolved: pytorch#858 The diff modifies the `padding` option and added tests with `compile`: * For the scaled_mm of shape MxKxN, the current `inner_padding` option only pads the `K` dimension. However, if `N` is not divisible by 16, we also got the error ``` E RuntimeError: CUDA error: CUBLAS_STATUS_NOT_SUPPORTED when calling `cublasLtMatmulAlgoGetHeuristic( ltHandle, computeDesc.descriptor(), Adesc.descriptor(), Bdesc.descriptor(), Cdesc.descriptor(), Ddesc.descriptor(), preference.descriptor(), 1, &heuristicResult, &returnedResult)` ``` So, modified the pad_inner option to also pad the K dimensions. ----- * The compile of inner-padding only works with the triton PR triton-lang/triton#4222. Before the triton PR, the inductor code-gen kernel fails at ``` tmp10 = tl.where(tmp6, tmp8, tmp9) TypeError: unexpected type fp8e5 and fp8e5 ``` Reviewed By: irobert0126 Differential Revision: D62003827

facebook-github-bot · 2024-09-16T07:46:50Z

This pull request was exported from Phabricator. Differential Revision: D62003827

vkuzo · 2024-09-16T20:03:21Z

test/float8/test_base.py

@@ -494,7 +494,7 @@ def test_different_configs_error(self):
        "base_dtype", [torch.float16, torch.bfloat16, torch.float32]
    )
    @pytest.mark.parametrize("use_fast_accum", [True, False])
-    def test_pad_inner_dim(self, base_dtype, use_fast_accum):
+    def test_pad_dimensions(self, base_dtype, use_fast_accum):


I think this would now need 3 cases:

pad N only

pad K only

pad N and K

vkuzo · 2024-09-16T20:03:46Z

test/float8/test_compile.py

@@ -40,14 +40,25 @@ def _test_compile_base(
    fullgraph: bool,
    config: Float8LinearConfig,
    dtype: torch.dtype,
+    pad_dimensions: bool,


same, I think we should test padding N, K, and N and K together

vkuzo · 2024-09-16T20:04:35Z

torchao/float8/config.py

@@ -120,7 +120,7 @@ class Float8LinearConfig:
    # inner dimension of a (dim 1) and b (dim 2) with 0s. This is needed for matmuls
    # _scaled_mm since it has the strong constraint that for M,N,K  N, K must be a multiple of 16.
    # This can cause a memory spike however so we keep this off by default.
-    pad_inner_dim: bool = False
+    pad_dimensions: bool = False


nit: update docblock

vkuzo · 2024-09-16T20:05:27Z

torchao/float8/float8_ops.py


+    return a_data, a_scale, b_data, b_scale, out_shape
+
+def postprocess_addmm(out: torch.Tensor, scaled_mm_config, out_shape):


nit: just inline instead of creating a new function? it's only two lines of code and used once

vkuzo · 2024-09-16T20:06:45Z

overall looks great! Can we have two more things:

update https://github.com/pytorch/ao/blob/main/benchmarks/float8/bench_padding.py and share data on how this PR impacts performance
mark this PR as BC-breaking and add the before and after state to the PR summary

y-sq · 2024-09-16T21:40:41Z

About the test failures:

Some aot_eager tests failed due to:

  E       RecursionError: maximum recursion depth exceeded while calling a Python object
  
  /opt/conda/envs/venv/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py:309: RecursionError
  _ test_aot_eager[dtype2-True-ScalingType.DYNAMIC-ScalingType.DELAYED-ScalingType.STATIC-True-True] _

From the error traces, it keeps calling the below functions for multiple times until exceeded the maximum recursion depth:

/opt/conda/envs/venv/lib/python3.9/site-packages/torch/_dynamo/backends/common.py:51: in _wrapped_bw_compiler
      return disable(disable(bw_compiler)(*args, **kwargs))
  /opt/conda/envs/venv/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py:632: in _fn
      return fn(*args, **kwargs)
  /opt/conda/envs/venv/lib/python3.9/site-packages/torch/_dynamo/backends/common.py:51: in _wrapped_bw_compiler
      return disable(disable(bw_compiler)(*args, **kwargs))
  /opt/conda/envs/venv/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py:632: in _fn
      return fn(*args, **kwargs)
  /opt/conda/envs/venv/lib/python3.9/site-packages/torch/_dynamo/backends/common.py:51: in _wrapped_bw_compiler
      return disable(disable(bw_compiler)(*args, **kwargs))
  /opt/conda/envs/venv/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py:632: in _fn

I'm still debugging this error.

test_inductor tests failed at

tmp10 = tl.where(tmp6, tmp8, tmp9)

TypeError: unexpected type fp8e5 and fp8e5

It's expected - the compile of inner-padding only works if the triton package contains this PR triton-lang/triton#4222.

How should we add such tests in github CI?

cc @vkuzo do you have any suggestions? thanks!

vkuzo · 2024-09-16T23:01:20Z

It's expected - the compile of inner-padding only works if the triton package contains this PR triton-lang/triton#4222

can you check if that PR is in triton 3.1.0? https://github.com/pytorch/pytorch/blob/main/.ci/docker/triton_version.txt is the current triton version in OSS PyTorch.

y-sq · 2024-09-16T23:10:58Z

@vkuzo , it's not included in triton 3.1. The pr (#4222) shows in the difference between 3.1.x and main: triton-lang/triton@release/3.1.x...main

drisspg

Yeah I think this is good, just also update the docs with the new behavior

Followup after pytorch/torchchat#815 to unblock migration to a newer version of PyTorch where AOTI seems to lost ability error out when one attempts to load CPU model on GPU, see https://github.com/pytorch/torchchat/actions/runs/9391753397/job/25913830802 for example Workaround by adding `-d ${DEVICE}` option to `aoti_runner`

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Sep 9, 2024

facebook-github-bot added the fb-exported label Sep 9, 2024

y-sq force-pushed the export-D62003827 branch from 0fdfe3e to 9481221 Compare September 10, 2024 21:06

y-sq force-pushed the export-D62003827 branch from 9481221 to 9ce26f8 Compare September 10, 2024 21:12

y-sq force-pushed the export-D62003827 branch from 9ce26f8 to 9b6212f Compare September 16, 2024 07:46

y-sq changed the title ~~Test compile with inner-padding~~ Some changes in inner-padding option Sep 16, 2024

vkuzo reviewed Sep 16, 2024

View reviewed changes

y-sq requested a review from drisspg October 2, 2024 05:59

drisspg approved these changes Oct 2, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Some changes in `inner-padding` option #858

Some changes in `inner-padding` option #858

y-sq commented Sep 9, 2024 •

edited

Loading

Uh oh!

pytorch-bot bot commented Sep 9, 2024 •

edited

Loading

Uh oh!

facebook-github-bot commented Sep 9, 2024

Uh oh!

facebook-github-bot commented Sep 10, 2024

Uh oh!

facebook-github-bot commented Sep 10, 2024

Uh oh!

facebook-github-bot commented Sep 16, 2024

Uh oh!

vkuzo Sep 16, 2024

Uh oh!

vkuzo Sep 16, 2024

Uh oh!

vkuzo Sep 16, 2024

Uh oh!

vkuzo Sep 16, 2024

Uh oh!

vkuzo commented Sep 16, 2024

Uh oh!

y-sq commented Sep 16, 2024 •

edited

Loading

Uh oh!

vkuzo commented Sep 16, 2024

Uh oh!

y-sq commented Sep 16, 2024

Uh oh!

drisspg left a comment

Uh oh!

Uh oh!


		return a_data, a_scale, b_data, b_scale, out_shape

		def postprocess_addmm(out: torch.Tensor, scaled_mm_config, out_shape):

Some changes in inner-padding option #858

Are you sure you want to change the base?

Some changes in inner-padding option #858

Conversation

y-sq commented Sep 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/858

❌ 1 New Failure

Uh oh!

facebook-github-bot commented Sep 9, 2024

Uh oh!

facebook-github-bot commented Sep 10, 2024

Uh oh!

facebook-github-bot commented Sep 10, 2024

Uh oh!

facebook-github-bot commented Sep 16, 2024

Uh oh!

vkuzo Sep 16, 2024

Choose a reason for hiding this comment

Uh oh!

vkuzo Sep 16, 2024

Choose a reason for hiding this comment

Uh oh!

vkuzo Sep 16, 2024

Choose a reason for hiding this comment

Uh oh!

vkuzo Sep 16, 2024

Choose a reason for hiding this comment

Uh oh!

vkuzo commented Sep 16, 2024

Uh oh!

y-sq commented Sep 16, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vkuzo commented Sep 16, 2024

Uh oh!

y-sq commented Sep 16, 2024

Uh oh!

drisspg left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Some changes in `inner-padding` option #858

Some changes in `inner-padding` option #858

y-sq commented Sep 9, 2024 •

edited

Loading

pytorch-bot bot commented Sep 9, 2024 •

edited

Loading

y-sq commented Sep 16, 2024 •

edited

Loading