feat: add support for fp32 kernels #56
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
pommedeterresautee
added
bug
…eir output tensor
|
opened an issue with reproduction code here: triton-lang/triton#674 |
# Conflicts: # optimizer/linear.py # test/test_linear_layer.py
|
measures done from this branch (with autocast) |
|
measures done on main (no autocast, full fp16) |
|
for memory, model in full fp16 in the autocast branch (so no autocast called): Shows that the code is as fast as in main when inference is not under autocast context manager. |
|
@gaetansnl there was a bug in Onnx Runtime, on main it's taking baseline model without setting fp16 to False, so it was working in full fp16 which doesn't work. In this branch there is no such flag and the Onnx model is in mixed precision |
|
with the weights in fp16 and the model in autocast |
gaetansnl
requested changes
Sep 23, 2022
| return out | ||
|
|
||
|
|
||
| def layer_norm(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, eps: float, implementation: JITFunction = _layer_norm_fwd_fused_single_pass): |
There was a problem hiding this comment.
why this one does not work as the other ? "output" is missing
There was a problem hiding this comment.
Maybe it's the opposite? Why attention has an output field? IMO it should be removed, just forgot.
For attention we create the output outside the function, and provide the tensor. The kernel is marked to convert all provided tensors to fp16 which includes the output in mixed precision. We should move it inside the function to avoid this unneeded casting.
Layernorm and linearlayer have not this issue by creating the output tensor of the right type from the begining.
There was a problem hiding this comment.
output is outside because it allows outside code to control allocations, not sure if it's still useful
| return outputs | ||
|
|
||
|
|
||
| def linear_layer(x: torch.Tensor, |
There was a problem hiding this comment.
same here ? why we don't have "output" ?
There was a problem hiding this comment.
see layernorm answer
test/test_attention.py
Outdated
|
|
||
| @pytest.mark.parametrize("batch", [1, 8, 32, 64]) | ||
| @pytest.mark.parametrize("implementation", ["torch", "triton_original", "triton"]) | ||
| def test_benchmark(benchmark, batch, implementation): | ||
| torch.manual_seed(0) |
There was a problem hiding this comment.
Should be moved at the beginning of the function to avoid mistakes IMO
There was a problem hiding this comment.
not sure to understand what you refer to?
There was a problem hiding this comment.
moved in an annotation
test/test_linear_layer.py
Outdated
| @@ -61,7 +53,7 @@ def test_benchmark(benchmark, shape: Shape, bias: bool, activation: str, contigu | |||
| else: | |||
| raise ValueError(f"Unknown activation: {activation}") | |||
|
|
|||
| torch_linear_layer = torch.nn.Linear(K, N, bias=bias, device="cuda", dtype=torch.float16) | |||
| torch_linear_layer = torch.nn.Linear(K, N, bias=bias, device="cuda", dtype=dtype) | |||
| torch_linear_layer.weight.data = layer_weight | |||
|
|
|||
| def torch_linear_activation(x): | |||
There was a problem hiding this comment.
same, you don't comapre to fp32 ?
There was a problem hiding this comment.
changed, ref is now fp32
test/test_torchdynamo_bert.py
Outdated
| (32, 16), (32, 128), (32, 256), | ||
| ], ids=lambda x: f"{x[0]}x{x[1]}") | ||
| @pytest.mark.parametrize("shape", [(bs, seq_l) for bs in [1, 8, 32] for seq_l in [16, 128, 256, 384, 512] | ||
| if bs * seq_l < 10000], ids=lambda x: f"{x[0]}x{x[1]}") | ||
| @pytest.mark.parametrize("implementation", implementations.keys()) | ||
| def test_benchmark_implementations(benchmark, model_reference_fp32, shape: (int, int), implementation: str): | ||
| torch.manual_seed(0) |
There was a problem hiding this comment.
why do we need the assert bellow
There was a problem hiding this comment.
which assert?
There was a problem hiding this comment.
assert implementation in implementations, f"unknown implementation: {implementation}"
There was a problem hiding this comment.
it has been removed
|
FYI, tests pass |
gaetansnl
reviewed
Sep 27, 2022
| return attention_forward_original(*args, **kwargs) | ||
|
|
||
|
|
||
| implementations = { |
There was a problem hiding this comment.
@pommedeterresautee Maybe we could use other config style to remove this from global scope https://docs.pytest.org/en/6.2.x/example/parametrize.html#paramexamples It makes things really hard to read IMO. And it will be worse as we add test
There was a problem hiding this comment.
Not sure if it's harder to read (it's local and there is the same pattern in all but one test not yet refactored, batched matmul), but I share your point about the fact that we want to control the number of implementations to test (probably at least have light and full flags), and doing it through the command line using pytest is certainly the best way. It is also true for number of shapes/batch sizes to test.
For that reason, that part of the code would be moved outside the global context, but it's not clear for me what it should look like.
It makes me think that it should be done in a dedicated PR, are you ok with that? If ok I write the issue.
gaetansnl
approved these changes
Sep 28, 2022
This was
linked to
issues
Sep 28, 2022
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
fix #39
fix #44
behavior of autocast: https://h-huang.github.io/tutorials/advanced/dispatcher.html#autocast + https://pytorch.org/docs/stable/amp.html