GitLab

* Enforce torch 2.0 and run attn tests with torch.compile
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* replace torch.compile with jit_fuser
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Added parallel cross entropy loss implementation using online softmax
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

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

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



* Added tests
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

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

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



* Added reshape of loss output
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

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

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



* Added to test list
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

* Added Triton dependency
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

* Added copyright
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

* Fixed lint errors
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

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

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



* Update setup.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Selvaraj Anandaraj <anandaraj@wisc.edu>

* Fixed lint and triton failure
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

* Removed flattening for scalars
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

* Skip tests on Blackwell due to TE CI caveat
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

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

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



* Added reason arg
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

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

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



* Do not register Triton dependency with setuptools
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>
Signed-off-by: Selvaraj Anandaraj <anandaraj@wisc.edu>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* non-exit tests
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* fixes L1 test

* fix test_multigpu_encoder

* fixes for other multi-encoder tests

* jax.extend.ffi to jax.ffi

* initialization with float32

* add init_dtype as an optional arg to all modules

* update use_scan query from xla flags

* relax threshold for test_encoder fp8

* relax the tols

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

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* add test_multiprocessing_encoder with processing spawning in bash

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

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

* add swa (left,0) + padding + brcm support
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



* final fixes
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



* upgrade to FE 1.9-rc
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* skip thd + CP + fused attn tests for cuDNN 9.6+ due to different stats shapes
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>

* Move test distributed encoder to L0 distributed test suit

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: Reese Wang <rewang@nvidia.com>

* draft implementation of fsdp2 fp8 all gather
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

* fix the convergence issue
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

* Add warning
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

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

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



* disable lint error
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

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

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



* fix the lint error
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

* fix lint error
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

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

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



* fix lint error
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

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

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



* fix lint error
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

* add comments
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

* add ref
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

* add related tests
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

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

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



---------
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Debug Mcore integration test

Avoid FP8 on Ampere and older. Generate synthetic data instead of depending on external data.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Handle deprecated `hidden_size` arg in norm modules
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Support initializing norm ops on CPU
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add integration test for Megatron-LM
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



* Rename Mcore integration test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Handle case in RMSNorm where hidden dim is not provided
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Implement ring attention primative for Jax.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Michael Goldfarb <mgoldfarb@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Add back the xla deterministic flag
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

* Add Userbuffers support for column TP linear layer
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add Userbuffers support for row TP linear layer
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Interpret linear+RS as row TP linear
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add Userbuffers support for FP8 row TP linear layer

Assumes FP8 RS, which is not a good assumption.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug bug with incorrect bias pointers in UB GEMM

Bias pointers are not properly offset for different data chunks. Also removed logic for FP8 RS.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add Userbuffers support for linear dgrad

Test passes with row TP, fails with col TP.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add Userbuffers support for linear wgrad
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add support for grad bias
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fused cast-transpose-dbias
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Support case where wgrad is optional
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Expand documentation
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



* Fix linter warnings
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use recently added convenience functions in Float8Tensor
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Respect autograd dtype
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



* Fix missing imports
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Respect PyT autocast dtype in bprop
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix linter warnings
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug merge conflicts
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



---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Skip encoder tests on V100

* Fix mulitprocessing jax.distributed.init

* Remove XLA xla_gpu_deterministic_ops which causes segfault

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>

remove one FA version in the L3 test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Remove PyTorch L0 distributed test

Forgot to remove in #1255.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug wheel test for PaddlePaddle
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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

* Reorganize PyTorch L1 tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Move ONNX tests to L1
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Move FA version test to L3
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Limit parallel build jobs in FA version test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Upgrade pylint and first round formatting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Format and fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

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

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

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

* WIP: make FA2 optional
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



* WIP: fix logic
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



* minor fixes
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 tweak
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add L1 test to test all supported FA versions
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update version to 2.1.1 and trim L1 tests
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



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

* remove onnxruntime from L1 FA versions tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

---------
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>

* Build custom ORT ops before running ONNX tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove ONNX from context parallelism tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Export ONNX ops that do compute in FP32

Matches internal impl of TE kernels.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add build script for custom ORT ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* add qkv descales to FA3
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 sbhd shapes
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



* force the same dtype when comparing FA3 and cuDNN FP8
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Revert "force the same dtype when comparing FA3 and cuDNN FP8"

This reverts commit 19e7f877026a19a32d2f02c6c9de20df4ae2e064.
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* force the same dtype when comparing FA3 and cuDNN FP8
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add try/except for FA3 when custom qkv descales are not supported
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace FA3 installation warning with a debug logging message
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



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

* avoid varlen_func for FP8 and improve messaging
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



* add SWA support for FA3
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* 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



* change preference reason for FP8 logic
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fixes
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
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

---------
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>

* Tests for distributed
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



* added the test to the qa script
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Changed qa
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix to test_numerics file
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* pr fixes
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



* fixes
Signed-off-by: Pawel Gadzinski <pgadzinski@pgadzinski-mlt.client.nvidia.com>

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

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



* Update tests/pytorch/distributed/run_numerics.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>

---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@pgadzinski-mlt.client.nvidia.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@pgadzinski-mlt.client.nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

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

* Improvements for wheels
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Fixes for wheel build
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Move package finder to common
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Lint
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>

* Fix CI and distributed test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* fix paddle ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* remove dtype from args
* update docs with permutation ops

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>

* Switch to nltk>3.8.1 and new data
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix nltk install
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

[C/PyTorch] Fixed incorrect use of `torch.distributed.new_group()` when creating intra-node group in `initialize_ub()` (#1087)

* updated initialize_ub() to use new_subgroups_by_enumeration() to generate intra-node groups, added new unit tests for TE layers with comm overlap
Signed-off-by: Alp Dener <adener@nvidia.com>

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

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



---------
Signed-off-by: Alp Dener <adener@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Fixes for wheels
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix paddle wheel test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Specify python version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add classifiers for python
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add utils to build wheels
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* make wheel scripts
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Fix paddle wheel
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* PaddlePaddle only builds for x86
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add optional fwk deps
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Python3.8; catch install error
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [wip] cudnn9 compile with paddle support
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [wip] dont link cudnn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

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

* dynamically load nvrtc
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



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

* remove residual packages; exclude stub from nvrtc .so search
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Exclude builtins from nvrtc .so search
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* properly include files for sdist
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* paddle wheel tie to python version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix paddle build from src [wip]
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix workflow paddle build
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



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

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

* fix lint from pr986
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Add sanity wheel test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add sanity import to wheel test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* remove upper limit on paddlepaddle version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Remove unused imports
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Remove pybind11 dependency
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix cpp tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Search .sos in cuda home
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



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

* CLeanup, remove residual code
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
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>

* Add enabled() to BasePrimitive

* Add layernorm/rmsnorm fallback

* Add cast_fp8 fallback

* Add transpose/cast_transpose XLA fall back

* Act_lu fallback

* Add transpose fallback

* Add softmax fallback

* Unify the use of _cast_fp8

* Add tests for NVTE_JAX_CUSTOM_CALLS_RE

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

* Add basic infrastructure for Sequential module
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add linear op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add FP8 support in linear op

Runs, but need to validate. Runtime errors with non-FP8 params and FP8 compute, or FP8 params and non-FP8 compute.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add reshape op and unit test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add bias op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add unfused linear op

Test does not pass with FP8.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug unfused linear op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add test for linear+bias op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add separate abstract classes for unfused and fused ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Consolidate unfused ops in submodule
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add linear-bias fused op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use fused cast-transpose in linear ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Disable GEMM+bias fusion with FP32 activations

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

* Add parallel unit test for unfused linear op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Refactor parallel tests to reduce job launches
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add all-reduce, all-gather, and reduce-scatter ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove unused file
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug multi-GPU FP8 test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add support for FP8 scale updates

Still need to implement amax reductions.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add license boilerplate
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fuse GEMM+bias in row TP

Add documentation for unfused ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Rename pipeline to fuser

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

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

* Preserve cached FP8 transpose between ops
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add option for fused wgrad accumulation
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Directly output FP8 from linear if needed
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix cuDNN front-end commit
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use updated FP8 tensor API for transpose caching
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use updated API for FP8 scale updates
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add tests for non-default FP8 recipes
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Rename UnfusedOperation to BasicOperation
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add unit test to check amax reduction with fusable op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Operator autograd state no longer needs to be initialized
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Initial functional implementation of linear op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug fused linear+bias op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove autograd context from functional linear impl
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use functional linear impl in fused linear+bias op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Rename subdirectory from "fuser" to "ops"

Avoid confusion with kernel fusers and graph compilers.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update with Float8Tensor changes in #820
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove unnecessary CPU overheads
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Correctly pass FP8 metadata from next op
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix linter errors
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add convenience functions to manipulate Sequential class
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update name of PyTorch extensions module
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Clear saved tensor data in linear op after bprop
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix Pylint error
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update name of PyTorch extensions module
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix test name in QA script
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update name of PyTorch extensions module
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Run distributed tests even when only 1 GPU is available
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Only run distributed tests with 2 GPUs if there are >=2 GPUs
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



* Review suggestions from @sudhakarsingh27 and @ksivaman

Fix spelling of "fusible". Avoid "input" name in internal APIs.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update transformer_engine/pytorch/ops/__init__.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Add example of formatting files locally
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* 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>

* 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>

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