Skip to content
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

Hierarchical CP implementation (Ulysses + Ring) #1209

Merged
merged 35 commits into from
Oct 7, 2024
Merged

Hierarchical CP implementation (Ulysses + Ring) #1209

merged 35 commits into from
Oct 7, 2024

Conversation

xrennvidia
Copy link
Collaborator

@xrennvidia xrennvidia commented Sep 27, 2024 •
edited by cyanguwa
Loading

Description

This PR adds a hierarchical implementation of context parallelism to attention. It uses A2A communications in low-level CP groups (e.g., via NVLink), and P2P communications in high-level CP groups (e.g., via IBLink). For more details, please refer to LongVILA and USP.

This implementation supports:

  • backends: FusedAttention, FlashAttention
  • dtype: BF16/FP16 (both backends), FP8 (FusedAttention only)
  • mask_type: causal, no_mask
  • attention types: MHA, MQA/GQA
  • qkv_format: sbhd, bshd
  • bias type: no_bias

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refractor

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

xrennvidia and others added 22 commits September 19, 2024 14:44
@xrennvidia
change API for hierarchical CP
2ee1389 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
Merge branch 'main' into xren/cp_a2a_p2p
3743c5f
@xrennvidia
move fp8 code before qkv reshape
0e826e7 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
try to insert A2A for hierarchical CP
44db43b 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
Merge branch 'main' into xren/cp_a2a_p2p
e717b81
@xrennvidia
make fwd work
705da7e 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
remove a redundant sync
99ace46 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
make bwd of hierarchical CP work
f73b2fb 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
fix dout a2a in bwd
0be5dd7 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
fix q_f16 with fp8
a6833c7 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
assert hierarchical CP implementation does not support THD format
caf3746 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
bug fix
03806c4 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
Merge branch 'main' into xren/cp_a2a_p2p
d3d336a
@xrennvidia
assert hierarchical CP does not support attn bias
ea1e4a3 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
add unit test for hierarchical CP
3cd21ab 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
fix cp_comm_type in unit test
90c0bb8 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
bug fix and code cleaning
8c6139d 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
Merge branch 'main' into xren/cp_a2a_p2p
52c7ae3
@xrennvidia
minor change
1db4e8b 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
an assert info change
1655edc 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
dout shape fix
edbe898 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
b096051 
for more information, see https://pre-commit.ci
@xrennvidia xrennvidia changed the title Hierarchical CP implementation (Ulysses + RIing) Hierarchical CP implementation (Ulysses + Ring) Sep 27, 2024
@xrennvidia
Copy link
Collaborator Author

/te-ci pytorch

@xrennvidia
Copy link
Collaborator Author

/te-ci pytorch

@xrennvidia
move function definitions to the front of the first call
67c590d 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
Copy link
Collaborator Author

/te-ci pytorch

cyanguwa
cyanguwa reviewed Oct 1, 2024
View reviewed changes
tests/pytorch/fused_attn/run_fused_attn_with_cp.py Outdated Show resolved Hide resolved
tests/pytorch/fused_attn/run_fused_attn_with_cp.py Show resolved Hide resolved
tests/pytorch/fused_attn/test_fused_attn_with_cp.py Outdated Show resolved Hide resolved
transformer_engine/pytorch/attention.py Outdated Show resolved Hide resolved
transformer_engine/pytorch/attention.py Show resolved Hide resolved
transformer_engine/pytorch/attention.py Outdated Show resolved Hide resolved
transformer_engine/pytorch/attention.py Show resolved Hide resolved
transformer_engine/pytorch/attention.py Outdated Show resolved Hide resolved
@cyanguwa
Copy link
Collaborator

cyanguwa commented Oct 1, 2024

@xrennvidia thanks for the PR! I left a few comments and also edited the PR description a bit. Let me know if it's accurate. Thanks!

xrennvidia and others added 9 commits October 1, 2024 15:19
@xrennvidia
fix tensor view comments
ba8844b 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
refine CP unit test
2aeb04a 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
512a771 
for more information, see https://pre-commit.ci
@xrennvidia
typo fix
d69792b 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
Merge branch 'xren/cp_a2a_p2p' of github.com:xrennvidia/TransformerEn…
fb8094a 
…gine into xren/cp_a2a_p2p
@xrennvidia
typo fix
caeddbf 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
save cp_size_a2a and rank_a2a in fwd
aa7e032 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
add more explainations of cp_group in doc_string
582a4e0 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
Merge branch 'main' into xren/cp_a2a_p2p
071c121
@xrennvidia
Copy link
Collaborator Author

@xrennvidia thanks for the PR! I left a few comments and also edited the PR description a bit. Let me know if it's accurate. Thanks!

It's accurate, thanks.

@xrennvidia
Copy link
Collaborator Author

/te-ci pytorch

cyanguwa
cyanguwa approved these changes Oct 7, 2024
View reviewed changes
Copy link
Collaborator

@cyanguwa cyanguwa left a comment

Choose a reason for hiding this comment

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

LGTM

@cyanguwa cyanguwa merged commit c24a4c4 into NVIDIA:main Oct 7, 2024
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@cyanguwa cyanguwa cyanguwa approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@xrennvidia @cyanguwa