Skip to content

[Paddle] Eliminate amax update bubbles by using custom_ops - #436

Merged
timmoon10 merged 2 commits into
NVIDIA:mainfrom
zlsh80826:rewang/eliminate-amax-update-bubbles
Sep 22, 2023
Merged

[Paddle] Eliminate amax update bubbles by using custom_ops#436
timmoon10 merged 2 commits into
NVIDIA:mainfrom
zlsh80826:rewang/eliminate-amax-update-bubbles

Conversation

@zlsh80826

@zlsh80826zlsh80826 commented Sep 18, 2023

Copy link
Copy Markdown
Collaborator

Paddle dynamic graph inserted a few bubbles when using assign operators/slice operators.
This PR moves amax_and_scale_update and copy_amax_from_buffer to custom_ops in c level to eliminate those tiny kernels and bubbles.

  • copy_amax_from_buffer: There was a slice kernel for accessing amax_history[0] and multiple kernels for = operator. By using the customized copy_amax_from_buffer op, all of above kernels can be reduced as a single memcpy. The spent time reduced from 18us to 1.8us.

  • amax_and_scale_update: There were multiple tiny kernels for updating amax/scale/scale_inv/history. This PR fused a few assign/elementwise operators to a single kernel to reduce the total 13 kernels to 4 kernels. The spent time reduced from ~39us to ~15us.

@zlsh80826
zlsh80826force-pushed the rewang/eliminate-amax-update-bubbles branch from cdd2a51 to 4d3b5a8CompareSeptember 18, 2023 15:48
@zlsh80826

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@zlsh80826
zlsh80826force-pushed the rewang/eliminate-amax-update-bubbles branch from 4d3b5a8 to 81fcae4CompareSeptember 19, 2023 06:50
Signed-off-by: rewang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed the rewang/eliminate-amax-update-bubbles branch from 81fcae4 to 7b55091CompareSeptember 19, 2023 06:52
@zlsh80826
zlsh80826 marked this pull request as ready for review September 19, 2023 07:01
@zlsh80826

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

Comment threadtransformer_engine/paddle/csrc/custom_ops.cu Outdated
Signed-off-by: rewang <rewang@nvidia.com>
@zlsh80826

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@timmoon10timmoon10 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

LGTM

@timmoon10
timmoon10 merged commit c32a62c into NVIDIA:mainSep 22, 2023
ptrendx pushed a commit that referenced this pull request Oct 2, 2023
* Eliminate amax_and_scale_update bubbles
Signed-off-by: rewang <rewang@nvidia.com>
* Add CUDA check
Signed-off-by: rewang <rewang@nvidia.com>
---------
Signed-off-by: rewang <rewang@nvidia.com>
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@zlsh80826@timmoon10@jeng1220