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

Ifu20240625 group gemm yewang12 #59

Open
wants to merge 25 commits into
base: dev
Choose a base branch
from
Open

Ifu20240625 group gemm yewang12 #59

wants to merge 25 commits into from
+38,958 −36,083

Conversation

wangye805
Copy link
Contributor

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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

Changes

Please list the changes introduced in this PR:

  • a4e95e8: the first commit for grouped linear from our upstream
  • Change B

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

cyanguwa and others added 25 commits June 13, 2024 21:25
@cyanguwa
Add documentation for dot product attention (#889)
4356938 
* add attention docs

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* WIP: update attention doc

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* WIP: update attention doc

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* WIP: update attention doc

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* WIP: update attn doc

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* WIP: update attn doc

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* WIP: update attn doc

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* WIP: update attention doc

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* first draft

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor tweak to first draft

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up pictures

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* first draft for review

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fixes

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add logging info/debug

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fix of an SWA message

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* use subprocess instaed of os.sys

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up benchmark script

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add example script and update notebook

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor tweak

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor tweaks

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix lint

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix Jax/Paddle related comments

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* rerun H100 benchmark

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* restrict fp8 tests to sm90+

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* move get_cudnn_version from common to pytorch utils

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

---------

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@ksivaman
Add auto-formatter (#919)
d99142a 
* Initial config test

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* remove linters, fix clang-format

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix clang-format

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix clang-format

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Remove lint

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Adjust config

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* use config file

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* adjust pylintrc

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* pre-format fixes

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Python only

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add FA module

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fixes

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update CI configs

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* CRLF -> LF

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* format

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* revert accidental formatting changes

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* try with sudo

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* cpp formatting

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix pylint error properly

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* some review comments

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* lint fixes

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* add fp8 attn include in the correct file

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* autofix PRs

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Apply formatting (#929)
9416519 
* Apply formatting

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Apply formatting

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@shamisp @ksivaman @pre-commit-ci
A hot fix to disable CE deadlock check (#926)
d71fc94 
* A hot fix to disable CE deadlock check

Signed-off-by: Pavel Shamis (Pasha) <pasharesearch@gmail.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Pavel Shamis (Pasha) <pasharesearch@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@cyanguwa @pre-commit-ci
[PyTorch] Adjust checkpointing of FP8 metadata for attention (#917)
fe5aa60 
* subclass DPA with BaseModule and test with test_gpt_checkpointing

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* test DPA only

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* test save and load

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove debug info

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor tweaks

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor tweak

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add hook in case core_attention._extra_state is missing

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* check named buffers in BaseModule; remove FP8 scratchpad override function; test FP8 for sm90+

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* minor fixes: test size, interval in recipe, named_buffer loop

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* move BaseModule from FusedAttention to DPA

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@phu0ngng
[Common] Remove CheckTensor if the workspace is empty in cast_transpo…
4a4f05d 
…se_fused (#931)

* rm tensor check if the workspace is empty

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* add trust_remote=true for load_dataset() in the mnist test

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@denera
[JAX] Fixing unused-variable warning at TE/JAX extension compile (#…
e2caf78 
…937)

replaced plain C asserts with NVTE_CHECK to avoid unused-variable warnings

Signed-off-by: Alp Dener <adener@nvidia.com>
@erhoo82 @timmoon10 @pre-commit-ci
Add the option to use SM for P2P comm in TP overlap (#914)
f458fcf 
* Add the option to use SM for P2P comm in TP overlap

Signed-off-by: Sangkug Lym <slym@nvidia.com>

* cleanup

Signed-off-by: Sangkug Lym <slym@nvidia.com>

* Python formatting with black

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Format C++ with clang-format

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update transformer_engine/pytorch/csrc/comm_gemm_overlap.h

Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

---------

Signed-off-by: Sangkug Lym <slym@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@ksivaman @pre-commit-ci
Remove leftover implementations for optional userbuffers support (#932)
16c8257 
* Remove optional UB build leftovers

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* rm unused import

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@ksivaman
Changed version to 1.9.0.dev
29e8bfc 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwa
[PyTorch] Fix tp_group_initialized error (#939)
94a426b 
fix tp_initialized error

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa @pre-commit-ci
[C/PyTorch] Simplify THD offset tensors (#927)
70d3251 
* simplify offset tensors

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fixes; tests pass

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix C lint

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace with_offset with with_padding

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace with_padding with padded

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* minor fixes after merge

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* minor fix for fused attn fwd/bwd calls

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix Jax

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* adjust spacing in docstring

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix pytorch tests; fix paddle api

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix lint

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix attn_biases

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix AttnFuncWithCP backward

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix jax

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix attn with CP

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix paddle

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@timmoon10
[PyTorch] Release GIL in PyTorch extensions (#938)
6ee92c4 
Release GIL in PyTorch pybind11 functions

Signed-off-by: Tim Moon <tmoon@nvidia.com>
@phu0ngng
Improve JAX build tool (#942)
85aeb90 
* adding option to select only .cpp files in a dir in the build tool

* change cmake build path

---------

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@yaox12 @StudyingShao
@QiZhangNV @phu0ngng
[Common/PyTorch] Grouped GEMM via multi-stream cuBLAS (#853)
a4e95e8 
* GroupedGEMM via multi-stream cublas

* fix A/B is nullptr while D is not nullptr

* add fp8 grouped gemm

* register with TorchScript

* add the GroupedLinear layer

---------

Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: Jiang Shao <jiangs@nvidia.com>
Co-authored-by: Qi Zhang <qizhang@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
@wenchenvincent
Merge commit 'a4e95e8' into ifu20240625_group_gemm (Conflicts
92a92d0 
unresolved).
NVIDIA/TransformerEngine@a4e95e8
@wenchenvincent @wangye805
Resolved conflicts for Merge commit 'a4e95e8' into ifu20240625_group_…
6a24374 
…gemm
@wangye805
[ROCm] skip the fp8 fused attn tests and rocm fused attn cuda graph t…
577e2ff 
…ests
@ipanfilo @wangye805
Fixed wrong noop tensors concatenation if data_storage is not big enough
f0f020f
@ipanfilo @wangye805
Address PR comments. Correct device capability checking for AMD GPUs.…
2a48ebf 
… Ignore MVTE_FLASH_ATTN env till FA is enabled for ROCm
@ipanfilo @wangye805
Remove unused variable
c0a4dd1
@timmoon10 @wangye805
[Core] Fix bug when selecting tuned RMSNorm kernels (#983)
191cd6f 
Fix typo when selecting tuned RMSNorm kernels

Signed-off-by: Tim Moon <tmoon@nvidia.com>
@ipanfilo @wangye805
Fix JAX examples, fix ROCm device capability check
abdbe31
@wangye805
[ROCm] resolve jax conflicts after IFU to grouped linear features
dfc87ea
@wangye805
[ROCm] merge the dev branch after IFU20240613
bf18926
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ipanfilo ipanfilo Awaiting requested review from ipanfilo

@wenchenvincent wenchenvincent Awaiting requested review from wenchenvincent

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

Successfully merging this pull request may close these issues.