Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Update megatron fused softmax follow megatron-lm#1539

Merged
crcrpar merged 7 commits into
NVIDIA:masterfrom
yaoyu-33:master
Nov 22, 2022
Merged

Update megatron fused softmax follow megatron-lm#1539
crcrpar merged 7 commits into
NVIDIA:masterfrom
yaoyu-33:master

Conversation

@yaoyu-33
Copy link
Copy Markdown
Contributor

We updated megatron fused softmax in the following aspects:

  1. We updated the limit of sequence length from 2048 to 4096 according to https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_softmax.py#L171
  2. We also enabled mask=None support in scaled_masked_softmax according to https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_softmax.py#L84

Copy link
Copy Markdown
Collaborator

@crcrpar crcrpar left a comment

Choose a reason for hiding this comment

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

how is the custom C++/CUDA extension of scaled_softmax_cuda built?
usually setup.py needs to be updated.
also, could you update https://github.com/NVIDIA/apex/blob/master/tests/L0/run_transformer/test_fused_softmax.py as well?

@yaoyu-33
Copy link
Copy Markdown
Contributor Author

@crcrpar comments addressed

Comment on lines +187 to +188
expected.backward(g0)
actual.backward(g1)
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

feels like these lines are based off of some existing cases but how about either checking grads or remove these lines?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

added grad check

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I think the purpose in other tests was just to test backward can work correctly

Copy link
Copy Markdown
Collaborator

@crcrpar crcrpar left a comment

Choose a reason for hiding this comment

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

generally okay but could you maybe tweak tests a bit?

Comment thread tests/L0/run_transformer/test_fused_softmax.py Outdated
Comment thread tests/L0/run_transformer/test_fused_softmax.py Outdated
@crcrpar crcrpar merged commit abeca58 into NVIDIA:master Nov 22, 2022
crcrpar pushed a commit to crcrpar/apex that referenced this pull request Nov 22, 2022
* Update megatron fused softmax follow megatron-lm

Signed-off-by: Yu Yao <[email protected]>

* Add mask=None support in scaled_masked_softmax

Signed-off-by: Yu Yao <[email protected]>

* Update setup.py for scaled_softmax_cuda

Signed-off-by: Yu Yao <[email protected]>

* Add tests for fused_scale_softmax (mask=None)

Signed-off-by: Yu Yao <[email protected]>

* Assert grad equal in fused softmax test

Signed-off-by: Yu Yao <[email protected]>

* Revert "Assert grad equal in fused softmax test"

Signed-off-by: Yu Yao <[email protected]>

Signed-off-by: Yu Yao <[email protected]>
Co-authored-by: Yu Yao <[email protected]>
hubertlu-tw pushed a commit to ROCm/apex that referenced this pull request Dec 29, 2022
* Update megatron fused softmax follow megatron-lm

Signed-off-by: Yu Yao <[email protected]>

* Add mask=None support in scaled_masked_softmax

Signed-off-by: Yu Yao <[email protected]>

* Update setup.py for scaled_softmax_cuda

Signed-off-by: Yu Yao <[email protected]>

* Add tests for fused_scale_softmax (mask=None)

Signed-off-by: Yu Yao <[email protected]>

* Assert grad equal in fused softmax test

Signed-off-by: Yu Yao <[email protected]>

* Revert "Assert grad equal in fused softmax test"

Signed-off-by: Yu Yao <[email protected]>

Signed-off-by: Yu Yao <[email protected]>
Co-authored-by: Yu Yao <[email protected]>
yuanzhedong pushed a commit to yuanzhedong/apex that referenced this pull request Jul 14, 2023
* Update megatron fused softmax follow megatron-lm

Signed-off-by: Yu Yao <[email protected]>

* Add mask=None support in scaled_masked_softmax

Signed-off-by: Yu Yao <[email protected]>

* Update setup.py for scaled_softmax_cuda

Signed-off-by: Yu Yao <[email protected]>

* Add tests for fused_scale_softmax (mask=None)

Signed-off-by: Yu Yao <[email protected]>

* Assert grad equal in fused softmax test

Signed-off-by: Yu Yao <[email protected]>

* Revert "Assert grad equal in fused softmax test"

Signed-off-by: Yu Yao <[email protected]>

Signed-off-by: Yu Yao <[email protected]>
Co-authored-by: Yu Yao <[email protected]>
Sign up for free to 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.

2 participants