Skip to content

Optimize pytorch layer_norm forward #20345

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

Jump to bottom
Closed
xiaomengy wants to merge 1 commit into from
Closed

Optimize pytorch layer_norm forward #20345

xiaomengy wants to merge 1 commit into from

Conversation

@xiaomengy
Copy link
Contributor

@xiaomengy xiaomengy commented May 10, 2019

Summary:
Seperate from D15194600
Optimize pytorch layer_norm op part 1:
optimize layer_norm_forward_cpu
import Eigen Maps for the performance of reduction

Differential Revision: D15290608

soumith
soumith suggested changes May 10, 2019
View reviewed changes
Copy link
Contributor

@soumith soumith left a comment

Choose a reason for hiding this comment

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

this PR now introduces Eigen into pytorch optimizations -- without discussion.
The previous comments of moving to native/cpu and using dispatcher are not addressed

@xiaomengy xiaomengy requested a review from soumith May 15, 2019 17:23
@xiaomengy
Copy link
Contributor Author

xiaomengy commented May 15, 2019

Removed Eigen now. For reduce case, I found that even using Vec256 is still slower than Eigen. So maybe we can consider to discuss or figure out how to make it better.

VitalyFedyunin
VitalyFedyunin reviewed May 15, 2019
View reviewed changes
aten/src/ATen/native/layer_norm.cpp Outdated
Copy link
Contributor

@VitalyFedyunin VitalyFedyunin May 15, 2019

Choose a reason for hiding this comment

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

Please move LayerNormForwardCPUImpl and all CPU logic into the /cpu subfolder using DECLARE_DISPATCH logic (see CopyKernel for the reference). It will allow to utilize AVX2 instructions and other optimizations of OSS build.

Copy link
Contributor Author

@xiaomengy xiaomengy May 16, 2019

Choose a reason for hiding this comment

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

Done

VitalyFedyunin
VitalyFedyunin reviewed May 15, 2019
View reviewed changes
aten/src/ATen/native/layer_norm.cpp Outdated
Copy link
Contributor

@VitalyFedyunin VitalyFedyunin May 15, 2019

Choose a reason for hiding this comment

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

I'm quite sure you want to call at:native::empty_like here and avoid all tracing/profiling checks.

VitalyFedyunin
VitalyFedyunin suggested changes May 15, 2019
View reviewed changes
Copy link
Contributor

@VitalyFedyunin VitalyFedyunin left a comment

Choose a reason for hiding this comment

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

This PR requires new tests to be written for native_layer_norm

@cpuhrsch
Copy link
Contributor

cpuhrsch commented May 15, 2019

@BIT-silence - Is this a copy of the C2 kernels?

@VitalyFedyunin
Copy link
Contributor

VitalyFedyunin commented May 15, 2019

You are exposing native_layer_norm function to the user level api which seems to be requiring contiguous inputs. Please add tests to cover APi calls as well as contiguous checks. Alternatively if intend is to introduce fast kernel, implement all code as kernel and avoid creating a new native function.

Also it would be nice to see benchmarks that compares calling old and new layer_norm code.

@xiaomengy
Copy link
Contributor Author

xiaomengy commented May 15, 2019

@BIT-silence - Is this a copy of the C2 kernels?

Logically it is, while C2 kernel is implemented by using Eigen lib. I tested the performance difference, for the rowwise moments part, Eigen version will be faster than this version by using compiler's auto vectorization, elementwise affine part performs the same. While I also tried to use Vec256 in rowwise moments part, it actually a little slower than the for-loop with auto vectorization.

@pytorchbot pytorchbot added the module: cpu CPU specific problem (e.g., perf, algorithm) label May 16, 2019
@xiaomengy
Copy link
Contributor Author

xiaomengy commented May 16, 2019

Thanks for the advice. I have removed native_layer_norm function. Later I will add the backward part for layer_norm then do autograd for layer_norm instead of batch_norm.

Some benchmark result for this change.
input shape = [64, 128, 56, 56] and normalized_shape = [128, 56, 56] with elementwise_affine=True,
on devvm forward time from 350ms to 87.6ms. And with this approach, it can let hugging face BERT model forward about 12% faster.

@xiaomengy xiaomengy requested a review from VitalyFedyunin May 16, 2019 00:57
soumith
soumith approved these changes May 16, 2019
View reviewed changes
Copy link
Contributor

@soumith soumith left a comment

Choose a reason for hiding this comment

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

thanks for making the changes. From my side this is good to go.

@xiaomengy
Copy link
Contributor Author

xiaomengy commented May 16, 2019

This PR requires new tests to be written for native_layer_norm

Actually I have removed native_layer_norm. Currently we don't expose new functions and add a fast path for layer_norm. Is that fine?

@fmassa
Copy link
Member

fmassa commented May 16, 2019

We should ideally test that the fast path gives the same results as the slow path. But maybe this is implicitly tested in the cuda tests

@xiaomengy
Copy link
Contributor Author

xiaomengy commented May 16, 2019

We should ideally test that the fast path gives the same results as the slow path. But maybe this is implicitly tested in the cuda tests

Actually not only cuda, but also jit test covers that.

@xiaomengy
Copy link
Contributor Author

xiaomengy commented May 16, 2019

Since I will add the grad part for this fast path and then all the layer_norm on CPU should go through this path. As the current fast path is actually tested by existing tests, I'm wondering if there is any need to add some specific test for that.

@xiaomengy @facebook-github-bot
Optimize pytorch layer_norm forward (pytorch#20345)
b075610 
Summary:
Pull Request resolved: pytorch#20345

Seperate from D15194600
Optimize pytorch layer_norm op part 1:
  optimize layer_norm_forward_cpu
  import Eigen Maps for the performance of reduction

Reviewed By: zheng-xq

Differential Revision: D15290608

fbshipit-source-id: d5589f67c515644403ff3ad11006ec43bab18809
@facebook-github-bot
Copy link
Contributor

facebook-github-bot commented May 21, 2019

This pull request has been merged in c9da011.

zdevito pushed a commit to zdevito/ATen that referenced this pull request May 22, 2019
@xiaomengy @facebook-github-bot
Optimize pytorch layer_norm forward (#20345)
4c38d35 
Summary:
Pull Request resolved: pytorch/pytorch#20345

Seperate from D15194600
Optimize pytorch layer_norm op part 1:
  optimize layer_norm_forward_cpu
  import Eigen Maps for the performance of reduction

Reviewed By: zheng-xq

Differential Revision: D15290608

fbshipit-source-id: cf2c208dfd6fbcbc4c69db3ed60278d9bee156b5
@vkuzo vkuzo mentioned this pull request Mar 24, 2020
Closed
vkuzo added a commit that referenced this pull request Mar 24, 2020
@vkuzo
add quantized layer norm implementation
15d19f7 
Summary:

Adds a quantized implementation of LayerNorm for server.

Relevant PRs:
* #20345 (floating point LN)
* #33080 (quantized BN)

A future PR will add the Python wrapper.

Test Plan:

numerics match the floating point implementation
TODO: benchmarks

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Mar 24, 2020
@vkuzo
Update on "add quantized layer norm implementation"
bce59e9 
Summary:

Adds a quantized implementation of LayerNorm for server.

Relevant PRs:
* #20345 (floating point LN)
* #33080 (quantized BN)

A future PR will add the Python wrapper.

Test Plan:

numerics match the floating point implementation
TODO: benchmarks

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Mar 24, 2020
@vkuzo
Update on "add quantized layer norm implementation"
24c2bf5 
Summary:

Adds a quantized implementation of LayerNorm for server.

Relevant PRs:
* #20345 (floating point LN)
* #33080 (quantized BN)

A future PR will add the Python wrapper.

Test Plan:

numerics match the floating point implementation
TODO: benchmarks

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Mar 24, 2020
@vkuzo
add quantized layer norm implementation
040d962 
Summary:

Adds a quantized implementation of LayerNorm for server.

Relevant PRs:
* #20345 (floating point LN)
* #33080 (quantized BN)

A future PR will add the Python wrapper.

Test Plan:

numerics match the floating point implementation
TODO: benchmarks

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 3c3721f
Pull Request resolved: #35329
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@cpuhrsch cpuhrsch Awaiting requested review from cpuhrsch

@VitalyFedyunin VitalyFedyunin Awaiting requested review from VitalyFedyunin

1 more reviewer

@soumith soumith soumith approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

Merged module: cpu CPU specific problem (e.g., perf, algorithm)

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

8 participants

@xiaomengy @cpuhrsch @VitalyFedyunin @fmassa @facebook-github-bot @soumith @pytorchbot @mruberry