- 09 Aug, 2025 1 commit
-
-
Daniel Stokes authored
* fix: Add stream synchronization before destroying MPI communicator (#1979) Signed-off-by:
djns99 <40156487+djns99@users.noreply.github.com> * feat: Implement column-wise userbuffer overlap for comm+GEMM operations Add support for overlapping column-wise allgather communication with GEMM operations to improve training performance: * **Core infrastructure changes:** - Update bulk_overlap_columnwise_ag() to accept explicit stream parameter - Modify userbuffers send/recv loops to use rank-ordered iteration - Add userbuffers_send_all/recv_all function declarations * **Python integration:** - Add bulk_overlap_ag_with_external_gemm() C++ extension function - Expose new overlap function via pybind11 bindings - Update overlap method configurations to include more ring_exchange ops * **LayerNorm MLP optimization:** - Enable column-wise quantization for FC2 gradient output - Implement overlap of allgather communication with FC2 DGRAD GEMM - Use fill_userbuffers_buffer_for_all_gather for efficient buffering This optimization allows overlapping communication and computation phases more effectively, reducing training wall-clock time by hiding allgather latency behind GEMM execution. Signed-off-by:
Tim Moon <4406448+timmoon10@users.noreply.github.com>
-
- 08 Aug, 2025 3 commits
-
-
Phuong Nguyen authored
* enabled TE GEMM for all recipes Signed-off-by:
Phuong Nguyen <phuonguyen@nvidia.com> * add warnings Signed-off-by:
pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
-
Paweł Gadziński authored
* turn on userbuffers for layers without debug Signed-off-by:
Pawel Gadzinski <pgadzinski@nvidia.com> * code drop Signed-off-by:
Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-authored-by:
-
Phuong Nguyen authored
* rm cudaGraph compatible trait from GroupedGEMM and groupedQuantize Signed-off-by:
-
- 07 Aug, 2025 3 commits
-
-
Phuong Nguyen authored
* rm batch_dim, sequence_dim, sequence_parallel_output Signed-off-by:
-
ldl authored
[PyTorch] Fix numeric overflow caused by int-type parameters and return value in the roundup function (#2034) Signed-off-by:
lvdunlin <lvdunlin@xiaomi.com> Co-authored-by:
-
Xin Yao authored
* PDL for MXFP8 Quantize 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 --------- Signed-off-by:
-
- 06 Aug, 2025 5 commits
-
-
Phuong Nguyen authored
Revert "[JAX] Disable TE Norm Custom Calls (#1993)" This reverts commit 6c970612 . --------- Signed-off-by:
-
hx authored
[PyTorch] fix input_quantizer usage for save_original_input; fix blockwise FP8 convert_and_update_tensor (#1978) * fix input_quantizer in save_original_input bwd Signed-off-by:
Hongxiao Bai <hongxiaob@nvidia.com> * fix get shape of blockwise tensor with only compact colwise data Signed-off-by:
-
Xin Yao authored
* for loop Signed-off-by:
-
Phuong Nguyen authored
* remove dot1_output_axes Signed-off-by:
-
xiaoxi-wangfj authored
Signed-off-by:xiaoxi-wangfj <690912414@qq.com>
-
- 05 Aug, 2025 2 commits
-
-
Phuong Nguyen authored
Disable Norm custom calls Signed-off-by: -
Phuong Nguyen authored
* new gemm operand specs processing Signed-off-by:
-
- 04 Aug, 2025 2 commits
-
-
buptzyb authored
* fix graph static grad input reuse Signed-off-by:
Robin Zhang <robinz@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update Signed-off-by:
-
Tim Moon authored
* Add basic kernel for swapping first two tensor dims Signed-off-by:
Tim Moon <tmoon@nvidia.com> * Add NVRTC kernel for swapping first dims Signed-off-by:
-
- 01 Aug, 2025 3 commits
-
-
Shang Zhang authored
Fix tensor numel() return dtype The original dytpe int would be an issue if the tensor element numbers are larger than int32's range (which is not a super large number in real workloads) Signed-off-by:Shang Zhang <shangz@nvidia.com>
-
Hongbin Liu authored
* disable wgrad accumulation and reduce in backward() And manually launch it in backward_dw() Signed-off-by:
Hongbin Liu <hongbinl@nvidia.com> * format Signed-off-by:
-
Autumn1998 authored
* fix bug if all values<0 Signed-off-by:
tongliu <tongliu@nvidia.com> * minor fix Signed-off-by:
-
- 31 Jul, 2025 1 commit
-
-
Jan Bielak authored
* Compute amax in normalization kernels as long as the pointer is provided, even if using non quantized output Signed-off-by:
Jan Bielak <jbielak@nvidia.com> * Fuse amax computation into normalization forward Signed-off-by:
-
- 30 Jul, 2025 3 commits
-
-
jberchtold-nvidia authored
* TE primitive checkpointing policies Signed-off-by:
Jeremy Berchtold <jberchtold@nvidia.com> * Remove batched gemm policy Signed-off-by:
-
Jan Bielak authored
Refactor normalization.cpp to use quantizer logic introduced in #1952 instead of manual quantization Signed-off-by: -
Tim Moon authored
Fix merge conflict bug with clearing op outputs Signed-off-by:
-
- 29 Jul, 2025 4 commits
-
-
Jan Bielak authored
* Compute amax in activation kernels when the output pointer is provided, even for non-fp8 outputs Signed-off-by:
-
Jan Bielak authored
Signed-off-by:
-
Tim Moon authored
* remove reciprocal op Signed-off-by:
zhongboz <zhongboz@nvidia.com> * Refactor Quantizer::create_tensor function Signed-off-by:
-
Przemyslaw Tredak authored
Signed-off-by:Przemek Tredak <ptredak@nvidia.com>
-
- 25 Jul, 2025 5 commits
-
-
Tim Moon authored
* Add ops for dropout and constant scale Signed-off-by:
-
Selvaraj Anandaraj authored
* Fixed double buffering issue for assymetric layers Signed-off-by:
Selvaraj Anandaraj <selvaraja@login-ptyche01.ptyche.clusters.nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by:
-
Kirthi Shankar Sivamani authored
* Remove deprecated device arg Signed-off-by:
-
Evgeny Tsykunov authored
* Support RMSNorm for QK 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 * rms -> RMSNorm, l2 -> L2Normalization (align with current pattern) Signed-off-by:
-
buptzyb authored
* optimize static grad outputs Signed-off-by:
-
- 24 Jul, 2025 7 commits
-
-
Alp Dener authored
[JAX] Fixing GemmPrimitive partitioning rules to handle tensor-parallelism correctly for sequence-parallel inputs (#1980) * updated GemmPrimitive partitioning rules to explicitly control all-reduce vs. reduce-scatter for sequence-parallelism Signed-off-by:
Alp Dener <adener@nvidia.com> * corrected handling of FSDP sharding for the RHS operand Signed-off-by:
-
Oleg Goncharov authored
* Fixed integer overflow when computing offsets Signed-off-by:
Oleg Goncharov <ogoncharov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by:
-
Kshitij Lakhani authored
Fix cudnn versioning in support in PyTorch DPA and Fused attn Signed-off-by:Kshitij Janardan Lakhani <klakhani@nvidia.com>
-
Jan Bielak authored
* Mark output tensors as not deletable in backward Signed-off-by:
-
Kirthi Shankar Sivamani authored
Fix cuDNN lib runtime loading and simplify Signed-off-by: -
Evgeny Tsykunov authored
* Increase intermediate precision and reuse tensors from fwd Signed-off-by:
-
Phuong Nguyen authored
* add manage_primitives() helper * disable GEMM primitives for non-MXFP8 recipes * implement the NVTE_JAX_CUSTOM_CALLS + deprecate NVTE_JAX_CUSTOM_CALLS_RE * replace NVTE_JAX_CUSTOM_CALLS_RE with NVTE_JAX_CUSTOM_CALLS in TE tests and examples * fix use_jax_gemm contextmanager Signed-off-by:
-
- 23 Jul, 2025 1 commit
-
-
Charlene Yang authored
* fix current device for cuDNN/cuBLAS handles 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 unit test Signed-off-by:
-
