GitLab

* WIP: fp8 v1 fprop integration
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

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

* add more debug info
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fprop working for h1; w/ debug info
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* cleanup; bprop running but has mismatches
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add gitlab frontend as submodule
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up and add back v0.9.2 FE support; fprop/bprop passing with 5e-2 tols
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix after merge; add bias_b/h to caching descriptor
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* distinguish fwd/bwd tensor types for bprop
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fix for F16 cases; include added dqkv_type and d_scale_dp
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* adjust out shape for bwd in test
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add casting from/to FP8 to DPA module
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* WIP: support all sbhd/bshd layouts
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* add qkvpacked and kvpacked support in both FusedAttnFunc and C levels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove qkvpacked/kvpacked calls in DPA module (used for testing)
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* remove tp setup; add allow_non_contiguous; update FE; revert to sbh3d in tests; clean up
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add NVTE_FP8_DPA_BWD to control whether to use FP8 bwd or F16 bwd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix MQA/GQA in FP8 v1 API
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update FE to 705d8e3, with API change
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* restrict mha_fill for THD format
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix fused attn with CP and comment out is_alibi code
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* clean up FE0.9 vs FE1.0 FP8 implementations, and related unit tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* change NVTE_FP8_DPA_BWD default to 1, and fix its use in qkvpacked/kvpacked APIs
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix lint and self.tp_size/group in FusedAttention()
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* update FE to 6902c94
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* update to FE v1.3.0
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fixes for FP8 MHA with different configs
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* emit stats regardless of is_training
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix linear when input is not Float8Tensor
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix d_out type when f16 bprop
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* fix user buffer for layernorm_linear/linear and revert two FP8 casts in MHA
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add docstring for fp8_dpa/mha in recipe
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* fix backend selection to avoid FA
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace transpose with transpose_2d
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* use RMSE for FP8 unit tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* replace two more transpose with transpose_2d
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* add FP8 initialization to FusedAttention
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

* Revert "add FP8 initialization to FusedAttention"

This reverts commit 15fffd825d6f23f31ea709b16ba01dfd61efabf8.
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Change order of ctxs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* add back docs and mark as beta
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* minor fixes for tests and docs
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Tweak error message for invalid FP8 GEMM dims
Signed-off-by: Tim Moon <tmoon@nvidia.com>

[PyTorch] Workaround for incorrect output from torch.cuda.is_bf16_compatible() on V100s and TU102s (#626)

* replaced torch.cuda.is_bf16_compatible() with explicit sm_80 check via torch.cuda.get_device_capability()
Signed-off-by: Alp Dener <adener@nvidia.com>

* implement te.utils.is_bf16_compatible() to replace torch.cuda counterpart
Signed-off-by: Alp Dener <adener@nvidia.com>

---------
Signed-off-by: Alp Dener <adener@nvidia.com>

* Implemented deferred initialization via `device='meta'` option for te.Linear and added new PyTorch example to demonstrate its use with FullyShardedDataParallel execution.
Signed-off-by: Alp Dener <adener@nvidia.com>

* correcting Float8Tensor initialization and fixing linting errors
Signed-off-by: Alp Dener <adener@nvidia.com>

* removed duplicate code from upstream rebase, local tests passing
Signed-off-by: Alp Dener <adener@nvidia.com>

* improved comments/documentation for FSDP example
Signed-off-by: Alp Dener <adener@nvidia.com>

* converted reset_parameters() into a base module function
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed Float8Tensor creation with deferred init, all tests passing locally
Signed-off-by: Alp Dener <adener@nvidia.com>

* extended deferred initialization to all TE modules
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed linting errors
Signed-off-by: Alp Dener <adener@nvidia.com>

* removed unnecessary reference to the parent module of parameter, added clarifying comments in parameter reset
Signed-off-by: Alp Dener <adener@nvidia.com>

---------
Signed-off-by: Alp Dener <adener@nvidia.com>

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

* [PyTorch] Linear: fix computation for wgrad if sequence_parallel=True
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

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

* [PyTorch] Linear: fix line length
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>

* Simplify logic for saving input tensor for Linear backward
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

Improve PyTorch memory usage
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Debug PyTorch and Paddle tests on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Only run Paddle layer tests with cuDNN fMHA on supported archs
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Reduce JAX FP8 GEMM sizes

Avoid split-k kernels on Ada.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Disable JAX fused self-attention test on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Run supported fused attention tests on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Run supported fused attention JAX tests on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Enable Paddle fused attention on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update reference scale calculation in TensorFlow test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Restore backend support to reference FP8 attention impl in PyT test

Review suggestion from @cyanguwa
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Debug Paddle tests on Ada
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Loosen tolerances for Paddle attention tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Assume causal mask implies equal seqlens in Paddle attention tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Better dimension assert for FP8
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* Clearer error messages for dtype and shape assertions
Signed-off-by: Carlos Mocholí <carlossmocholi@gmail.com>

* Update transformer_engine/pytorch/utils.py
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Carlos Mocholí <carlossmocholi@gmail.com>

* Update transformer_engine/pytorch/utils.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Carlos Mocholí <carlossmocholi@gmail.com>

---------
Signed-off-by: Carlos Mocholí <carlossmocholi@gmail.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* Catch FP8 modulo16 error before cublas and fp8 kernels
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* ignore self attention mask for causal type
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* further relax checks to run FA, update docs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

* minimum ampere requirement for fa
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

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

retain grad related attrs while casting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>