GitLab

Reset FP8 weight workspace if usages are invalid
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* save original input
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

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

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



* fix input_quantizer usage in Linear bwd
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* minor fix
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* refine the docstring
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

* Merge remote-tracking branch 'origin/main' into save_bf16_in_fp8_gemm
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

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

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



* decouple linear bwd with save_original_input; clean up UTs
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>

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

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



---------
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>
Signed-off-by: hx <hongxiaob@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* some initial code
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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

* fixed returning layernorm etc
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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



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



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



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

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

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

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

* added pip install to test.sh
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Update transformer_engine/pytorch/export.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.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



* float8currentscaling quantizer exception
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 to wheels
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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

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

* fix
Signed-off-by: root <root@prenyx0221.a51.clusters.nvidia.com>

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

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



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



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

* fixes
Signed-off-by: root <pgadzinski@nvidia.com>

* fixes
Signed-off-by: root <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@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@nvidia.com>

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

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



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

* Update setup.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>

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

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



* onnxscript version chnage
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@nvidia.com>

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

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



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

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

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

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

---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>
Signed-off-by: root <root@prenyx0221.a51.clusters.nvidia.com>
Signed-off-by: root <pgadzinski@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@gmail.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: root <root@prenyx0221.a51.clusters.nvidia.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@gmail.com>

* Fix clearing tensor data in backward removing is_first_op
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Misc fixes
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Use Linear weight dtype and device for compute consistently
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Add backward dbias + quantize fusion
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Pass recipe to OperationFuser to allow recipe-dependent fusions
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Remove redundant view from activations
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Add bias activation backward fusion
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* Apply suggestions from code review
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Jan Bielak <jbielak@nvidia.com>

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

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



---------
Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Co-authored-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>

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

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

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Added MCore fsdp support for TE
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



* Refactored based on new MCore FSDP
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

* Code cleanup and extended across modules
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 FSDP suport
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>

---------
Signed-off-by: Selvaraj Anandaraj <selvaraja@cw-dfw-cs-001-login-01.cm.cluster>
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: Tim Moon <4406448+timmoon10@users.noreply.github.com>

Fix layernorm output shape in LayernormLinear, make it align with input shape
Signed-off-by: Guyue Huang <guyueh@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [PyTorch|common] Implement unpadding kernel for FP8

1. Add multi-tensor unpadding kernel
2. Replace split+cat with unpadding kernel in Fp8Padding and Fp8Unpadding
3. Add unpadding with padding unit tests
Signed-off-by: xiaoxi-wangfj <690912414@qq.com>

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

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



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

* Update padding.cu
Signed-off-by: Xin Yao <xiny@nvidia.com>

---------
Signed-off-by: xiaoxi-wangfj <690912414@qq.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>

* finish python ref impl for bulk alloc
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* c++ bulk alloc worked, still draft version
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* clean up
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* resolve rebase conflict
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* add license
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* use shared_ptr to auto manage reference count
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* attempt to fix misc training error
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* attempt to handle case where experts get zero token
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* updated with fused C++ function calls
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* clean up
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* experiment with reducing py object construction time
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix seg fault bug in inference mode
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix lint
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fuse torch split into bulk alloc
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* clean up
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* rebase to latest main
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix unit test failure
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix lint error
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* refactor create_tensor to use get_scale_shape
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* refactor quantize to call quantize_cpp
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Implement separate functions for multi-tensor quantize and split + multi-tensor quantize
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update grouped linear module with fused split+quantize func
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Move multi-tensor quantize func to cast.cpp
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Do not expose quantizer helper function externally
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix linter warnings
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



* Revert cuDNN frontend commit
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* fix corner cases with zero tokens
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* add comments
Signed-off-by: zhongboz <zhongboz@nvidia.com>

---------
Signed-off-by: zhongboz <zhongboz@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>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

* Do not initialize quantized weights with column-wise usage in inference mode
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Use no-grad mode instead of inference mode in tests
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
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: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add support for overlapping wgrad NCCL AG with dgrad GEMM
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

* Remove unused wait on memcpy API from UB
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

* Add better commenting to MXFP8 overlap
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>

---------
Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com>
Co-authored-by: dastokes <dastokes@dastokes-dvt-01.nvidia.com>

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

Fix for loading old ckpt formats
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

only get skip_fp8_weight_update in fp8_graph_capturing
Signed-off-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* refactor the multi_stream utils + implement nvte_multi_tensor_quantize in TE/Common

* implement GroupedQuantizer and grouped_quantize in jaxx

* fix logical_axes_names for transpose tensor in ScaledTensor
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: Hua Huang <huah@nvidia.com>
Co-authored-by: Ming Huang <mingh@nvidia.com>

[PyTorch] FP8 Subchannel Recipe With FP8 Gather And Configurable Scaling Factor Tensor Swizzling (#1707)

* functional kernel for columnwise + no-transpose option, still hacky
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* pass all quantizer unit tests
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* refactor, add gemm ready api
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* make format options private members, simplify api
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* swizzle scales right before gemm
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* bug fix of single layer test
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* attempt to fix lint issue
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* fp8 gather pass, need minor refine
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* fix return_layernorm_output_gathered case
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* remove special cases, add sanity check before gemm
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix lint
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* lint ungrouped imports
Signed-off-by: zhongboz <zhongboz@nvidia.com>

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

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



* Implement dequantize for compact 1D blocks.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* add more unit test with dequantize compact supported
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* lint again
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* make ag for subchannel respect async
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* zero tolerance in distributed test
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix zero tolerance test
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* resolve rebase issues
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* lint & format
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix lint
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* clean up
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* bug fix
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* relax rtol for fp32 distributed test
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix some ci issue
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix ci test failure in debug mode
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Force row-wise and column-wise data to have same data format

Prototype "all-gather usage" in quantizer.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove dead logic for high-precision AGs
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



* Debug FP8 block-wise tensor tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Handle case where LayerNormLinear returns gathered norm output
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* fix debug mode
Signed-off-by: zhongboz <zhongboz@nvidia.com>

---------
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: Keith Wyss <kwyss@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>
Co-authored-by: Keith Wyss <kwyss@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Quantizer update
Signed-off-by: Evgeny Tsykunov <etsykunov@etsykunov-mlt.client.nvidia.com>

* Update import
Signed-off-by: Evgeny <etsykunov@nvidia.com>

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

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



* Introduce _update_weight_quantizers and _get_weight_tensors/_get_weight_quantizers
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Add test
Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Move _quantizer to the QuantizedTensorBase
Signed-off-by: Evgeny <etsykunov@nvidia.com>

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

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



* Fix import
Signed-off-by: Evgeny Tsykunov <etsykunov@etsykunov-mlt.client.nvidia.com>

---------
Signed-off-by: Evgeny Tsykunov <etsykunov@etsykunov-mlt.client.nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Co-authored-by: Evgeny Tsykunov <etsykunov@etsykunov-mlt.client.nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>

* Use an empty torch tensor to indicate no fp8 information in extra_state
Signed-off-by: Peter St. John <pstjohn@nvidia.com>

* Add huggingface from_pretrained / save_pretrained tests

Adds integration tests to ensure models containing TransformerLayer
objects can be saved and loaded using the from_pretrained and
save_pretrained methods.
Signed-off-by: Peter St. John <pstjohn@nvidia.com>

---------
Signed-off-by: Peter St. John <pstjohn@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Check tensor-recipe compatibility
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

* Tensor class in recipe, checking for *Base
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

* Extend recipe __repr__ with recipe_type
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

* Warn about recipe change
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

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

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



* Enable dynamic recipe change: clear fp8 workspace
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

* TE 1.x checkpoint compatibility
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

* Disable warning for recipe wrappers
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

* Test recipe change
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

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

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



* Use QuantizedTensorBase
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

* Fix circular import
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

* Revert previous circular import fix
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

* Fix pytorch imports in common
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

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

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



* Let quantizer know about the recipe
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

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

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



* Fix imports
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

---------
Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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



* 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



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

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



* 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



* 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



* 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>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Przemek Tredak <ptredak@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Set sequence_parallel before super().__init__() in norm modules
Signed-off-by: Evgeny Tsykunov <etsykunov@etsykunov-mlt.client.nvidia.com>

* getattr(self, sequence_parallel, None) -> self.sequence_parallel
Signed-off-by: Evgeny Tsykunov <etsykunov@etsykunov-mlt.client.nvidia.com>

---------
Signed-off-by: Evgeny Tsykunov <etsykunov@etsykunov-mlt.client.nvidia.com>
Co-authored-by: Evgeny Tsykunov <etsykunov@etsykunov-mlt.client.nvidia.com>

* Avoid spurious warning with non-FP8 GroupedLinear
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use `QuantizedTensorBase`
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
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: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Initial work toward restoring UB support in te.Sequential
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Forward UB linear runs, but has numerical error
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

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

* Remove Python checks for MXFP8 UB linear forward
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add dim check for MXFP8 full tiles
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Move QuantizedTensor logic out of UB comm and into Python helper function
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Support MXFP8 AGs
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Coalesce NCCL all-gathers for MXFP8 all-gather
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Initial impl of backward UB linear in te.Sequential
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug UB linear backward with no quantization

dgrad GEMM + dx RS is still broken.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix chunk dims for dgrad GEMM + dx RS
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debugging MXFP8 UB cases

Still failing with dy AG + wgrad GEMM
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use NCCL to overlap dy AG with dgrad GEMM
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Initial refactoring of linear module forward
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Refactor linear module backward
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug linear module UB tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Do not store autograd context within wgrad GEMM closure
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix linter warnings
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 LayerNormLinear
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 LayerNormMLP
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Fix linter warnings
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



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

* Minor style tweaks
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 incorrect usage for GEMM input with block-scaled FP8
Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

* Disable dgrad GEMM + UB AG + NCCL AG overlapping
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



* Disable dgrad GEMM + UB AG + NCCL AG overlap in te.Sequential
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



* Restore support for internal quantized tensors
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add tests for MXFP8 GEMM with UB
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>

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

* Debug test failures
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>

Signed-off-by: Peter St. John <pstjohn@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Changes to Linear
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

* Removing unnecessary check
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Relax the absolute tolerance in FP32 distributed test
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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



* Add QuantizedTensorBase class
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Change the blockwise tensor.
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* A little cleaning
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

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

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



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

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

* Fix autocast deprecation warnings
Signed-off-by: Xin Yao <xiny@nvidia.com>

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

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

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

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Fix bug when reshaping norm output from LayerNormLinear
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* init
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@nvidia.com>

* clear() 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



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

* revert adding cpu offload tests for mxfp8
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



* 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



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



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

* forgot to uncomment assert
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Prevent using fp8 weights + non-fp8 computation
Signed-off-by: kunlunl <kunlunl@nvidia.com>

* Using warnings instead of raising an error
Signed-off-by: kunlunl <kunlunl@nvidia.com>

* Add dequantization back
Signed-off-by: kunlunl <kunlunl@nvidia.com>

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

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



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

[PyTorch] Fix cuBLAS workspace leak in applications that initialize+destroy Userbuffers more than once (#1715)

safeguarded cuBLAS workspace expansion in initialize_ub() to avoid exponential growth across repeat initializations
Signed-off-by: Alp Dener <adener@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* rtx5090 arch fix support
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* apprend `nvte` to the function name so that its visible in framework specific dirs
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix typo
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* add filter for nvte_is_supported_nontn_fp8_gemm
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* properly expose the api
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

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

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



* feedback from PR
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

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

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



* move the function to apt header/c files
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

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

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



* add more info
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

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

* split wgrad for GroupedLinear
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

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

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

Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* support wgrad split for linear and ln_linear
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* add comments and fix WeightGradStore
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* support bias and fix unit tests
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* minor fix
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* support fuse_grad_accumulation=false
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

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

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



* add wgrad split for layernorm_mlp
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* minor fix
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

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

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



* fix unittest
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

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

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



* add unittest for distributed interface apply Dener's suggestion
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

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

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



* minor fix
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* replace split_bw with delay_wgrad_compute
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

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

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



* Update transformer_engine/pytorch/module/layernorm_mlp.py
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update transformer_engine/pytorch/module/linear.py
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update transformer_engine/pytorch/module/layernorm_linear.py
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* remove comments
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

---------
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Hongbin Liu <hongbinl@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* re merge request
Signed-off-by: Peter Dykas <wdykas@nvidia.com>

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

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



* add docstring
Signed-off-by: Peter Dykas <wdykas@nvidia.com>

---------
Signed-off-by: Peter Dykas <wdykas@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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



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

---------
Signed-off-by: Pawel Gadzinski <pgadzinski@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>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

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

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

* file i forgot
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



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

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

* Update transformer_engine/debug/pytorch/utils.py
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>

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

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

* Update transformer_engine/pytorch/tensor/_internal/float8_tensor_base.py
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com>

* all tensor types
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

* 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



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

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



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

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

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

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

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

* added blockwise quantizer
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



* 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



* make debug option optional
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

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

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

* Add fp8 weight transpose cache check in backward, and regenerated it if it does not exist
Signed-off-by: jianbinc <shjwudp@gmail.com>

* Properly handle fsdp shard model weight input.
Signed-off-by: jianbinc <shjwudp@gmail.com>

* move Float8Tensor to QuantizedTensor in cast_master_weights_to_fp8 UT
Signed-off-by: jianbinc <shjwudp@gmail.com>

* handle Float8TensorBase issue
Signed-off-by: jianbinc <shjwudp@gmail.com>

* fix bug in activation recompute
Signed-off-by: jianbinc <shjwudp@gmail.com>

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

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



---------
Signed-off-by: jianbinc <shjwudp@gmail.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Enable MXFP8 and Per-Tensor Current Scaling for Grouped Linear
Signed-off-by: Xin Yao <xiny@nvidia.com>

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

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

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



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

* remove grouped linear parallel mode test
Signed-off-by: Xin Yao <xiny@nvidia.com>

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

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

* internal=False for now
Signed-off-by: Xin Yao <xiny@nvidia.com>

* remove unused import
Signed-off-by: Xin Yao <xiny@nvidia.com>

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

* Add GEMM logic for blockwise quantized tensors.

GEMM test cases included in pytorch integration.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update NVTE_BLOCK_SCALING for GEMM.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Gate feature on CUDA 12.9
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Gemm typo.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Remove unecessary type converter change.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Reflect epilogue availability and test supported epilogues.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* GEMM simplifications from recipe branch.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Format py code.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update GEMM DGelu tests to match support depending on output dtype.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Force pow2Scales in GEMM
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Add GEMM test to pytorch test suite.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Add copyright to GEMM test.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update import for GEMM test.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Add license.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update test gemm supported predicate.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Use sgemm like interfaces and naming.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Rewrite GEMM comment.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* MR Feedback.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Recipe setup for Linear modules.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Use 12.9 feature test.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Run against tensor dumps from internal library.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update FIXME to TODO with linked issue.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update full recompute feature to save recipe.

The recompute context uses the same recipe
and fp8 settings as the original fwd pass.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* MR Feedback. Avoid reusing quantizer objects.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update logic in module.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Format py.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update for PP bug.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update test numerics.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update force_power_of_2 scales in the recipe.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update usage method to satisfy upstream changes.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* fix subchannel recipe in distributed test with bf16 gather
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Edit and cleanup BF16 gather code.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update test import.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* support columnwise only mode to 1D quantize kernel
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Format and move enum
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Skip alloc.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* try async bf16 gather
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Format python code.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Document and type code.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update pytorch lint errors.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Dont set high precision dtype.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Add test for sanity and CG; fix CG for sequential?
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Keep make_quantizers API stable

Update num_quantizers instead to pass cuda_graph tests.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Fix import name.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Rename recipe method.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Skip grouped linear sanity test.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Set usage before BF16 gather.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* refactor for nvte_quantize_v2
Signed-off-by: zhongboz <zhongboz@nvidia.com>

* Format code.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Cleanup nvte_quantize_v2
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Test fp32 scales.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Disable CUDA graph.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Simplify layernorm linear
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Cleanup layernorm linear.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* LayerNorm linear bwd gather logic.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Communication updates.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update transformer_engine/pytorch/ops/op.py

Apply MR comment change.
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: kwyss-nvidia <kwyss@nvidia.com>

* Lint fix.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* MR feedback.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Enable cuda graph tests.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Reduce chance of spurious failure and reword.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Review suggestions from @timmoon10
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update CPP tests.
Signed-off-by: Keith Wyss <kwyss@nvidia.com>

* Update common.h
Signed-off-by: Xin Yao <yaox12@outlook.com>

* Update test_float8blockwisetensor.py
Signed-off-by: Xin Yao <yaox12@outlook.com>

---------
Signed-off-by: Keith Wyss <kwyss@nvidia.com>
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: kwyss-nvidia <kwyss@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Xin Yao <yaox12@outlook.com>
Co-authored-by: zhongboz <zhongboz@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-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>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Xin Yao <yaox12@outlook.com>

* Use dummy wgrads for lower memory consumption
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Bug fix to avoid sharing gradients.
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Disable automatic use of batch_p2p_comm for CP2
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

* Change weight to origin_weight for LN_LINEAR
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

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

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

Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Make sure that weight matrix has required usages for dgrad GEMM
Signed-off-by: Tim Moon <tmoon@nvidia.com>