Skip to content

Add attention bias and qkv format to context parallelism - #726

Merged
cyanguwa merged 26 commits into
NVIDIA:mainfrom
xrennvidia:xren/cp_with_attn_bias
Apr 26, 2024
Merged

Add attention bias and qkv format to context parallelism#726
cyanguwa merged 26 commits into
NVIDIA:mainfrom
xrennvidia:xren/cp_with_attn_bias

Conversation

@xrennvidia

@xrennvidiaxrennvidia commented Mar 19, 2024

Copy link
Copy Markdown
Collaborator

Add attn_bias and qkv_format to AttnFuncWithCP.

Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidiaxrennvidia changed the title [WIP] add attention bias to context parallelism[WIP] add attention bias and qkv format to context parallelismMar 21, 2024
@xrennvidiaxrennvidia changed the title [WIP] add attention bias and qkv format to context parallelismAdd attention bias and qkv format to context parallelismMar 21, 2024
@cyanguwa
cyanguwa self-requested a review March 26, 2024 18:56
Comment threadtests/pytorch/fused_attn/test_fused_attn_with_cp.py
Comment threadtransformer_engine/pytorch/attention.py
Comment threadtransformer_engine/pytorch/attention.py Outdated
@cyanguwa

Copy link
Copy Markdown
Collaborator

/te-ci pytorch

Comment threadtransformer_engine/pytorch/attention.py
@cyanguwa

Copy link
Copy Markdown
Collaborator

/te-ci pytorch

xrennvidiaand others added 2 commits April 24, 2024 16:26
@cyanguwa

Copy link
Copy Markdown
Collaborator

/te-ci pytorch

@cyanguwa
cyanguwa merged commit 9709147 into NVIDIA:mainApr 26, 2024
@xrennvidia
xrennvidia deleted the xren/cp_with_attn_bias branch April 26, 2024 18:28
phu0ngng referenced this pull request in phu0ngng/TransformerEngine May 3, 2024
* make FusedAttn with CP support bias
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* assert Alibi cannot work with CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* syntax fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* fix variable name
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* fix tensor shapes
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* a typo fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* fix bias indexing for CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* bug fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* add attn bias tests
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* change dbias update location
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* fix CP test model configs
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* change CP test sequence length
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* make AttnFuncWithCP support qkv format of sbhd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* make sure qkv are contiguous for CP in cuDNN fused attn
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* change assert message
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* fix code format
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
---------
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Co-authored-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>
pggPL pushed a commit to pggPL/TransformerEngine that referenced this pull request May 23, 2024
* make FusedAttn with CP support bias
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* assert Alibi cannot work with CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* syntax fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* fix variable name
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* fix tensor shapes
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* a typo fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* fix bias indexing for CP
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* bug fix
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* add attn bias tests
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* change dbias update location
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* fix CP test model configs
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* change CP test sequence length
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* make AttnFuncWithCP support qkv format of sbhd
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* make sure qkv are contiguous for CP in cuDNN fused attn
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* change assert message
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
* fix code format
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
---------
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Co-authored-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@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.

2 participants

@xrennvidia@cyanguwa