- 18 Nov, 2025 1 commit
-
-
Paweł Gadziński authored
* 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 packed versions Signed-off-by:
pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
-
- 14 Nov, 2025 1 commit
-
-
jberchtold-nvidia authored
* Refactor to avoid storing a global quantization config so direct recipe passing works as intended Signed-off-by:
Jeremy Berchtold <jberchtold@nvidia.com> * fix use_split_accumulator for current scaling recipe Signed-off-by:
greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by:
jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update transformer_engine/jax/quantize/quantizer.py Co-authored-by:
-
- 13 Nov, 2025 1 commit
-
-
jberchtold-nvidia authored
* Support for checkpointing quantizations Signed-off-by:
JAX Toolbox <jax@nvidia.com> * Checkpoint quantizations Signed-off-by:
-
- 10 Nov, 2025 1 commit
-
-
Teddy Do authored
* Changing default activations in MLP, TransformerLayer, dropout rate after FC1 to 0, and return_layernorm_output to False Signed-off-by:
tdophung <tdophung@nvidia.com> * Fixing the failing tests by hard coding arguments to the previous values instead of relying on newer default values Signed-off-by:
-
- 07 Nov, 2025 1 commit
-
-
jberchtold-nvidia authored
* Add test to check jaxpr that amax is reused for nvfp4 recipe Signed-off-by:
-
- 03 Nov, 2025 1 commit
-
-
Phuong Nguyen authored
* L1 rework Signed-off-by:
Phuong Nguyen <phuonguyen@nvidia.com> * comment out test_multi_process_grouped_gemm for now Signed-off-by:
-
- 30 Oct, 2025 1 commit
-
-
Kshitij Lakhani authored
* Fix: Skip determinism tests for bprop for all sm >=100 Signed-off-by:
Kshitij Lakhani <klakhani@nvidia.com> * Add username to TODO Signed-off-by:
-
- 23 Oct, 2025 1 commit
-
-
jberchtold-nvidia authored
* Make SR rng state always 2D (num_devices, 4) Signed-off-by:
-
- 22 Oct, 2025 1 commit
-
-
jberchtold-nvidia authored
* [JAX] Support recipe flags for disabling SR, RHT, and 2D quantization Signed-off-by:
-
- 17 Oct, 2025 1 commit
-
-
Kshitij Lakhani authored
* Fix imports in test for deprecated jax.experimental.pjit Signed-off-by:
Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by:
-
- 14 Oct, 2025 2 commits
-
-
Kirthi Shankar Sivamani authored
* Initial API change Signed-off-by:
Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Change all imports and api Signed-off-by:
Tim Moon <4406448+timmoon10@users.noreply.github.com> * fix nvfp4 test Signed-off-by:
-
Kshitij Lakhani authored
* Add BRCM support when creating a test mask for fused attn Signed-off-by:
-
- 09 Oct, 2025 1 commit
-
-
jberchtold-nvidia authored
Signed-off-by:
-
- 08 Oct, 2025 1 commit
-
-
Hua Huang authored
* Try async copy of grouped GEMM group_sizes data Signed-off-by:
Hua Huang <huah@nvidia.com> --------- Signed-off-by:
-
- 06 Oct, 2025 1 commit
-
-
Phuong Nguyen authored
* not fuse bias for output all reduction case + unit tests Signed-off-by:
-
- 03 Oct, 2025 1 commit
-
-
vthumbe1503 authored
Signed-off-by:Varun Thumbe <vthumbe@nvidia.com> *Jax integration for clamped swiglu. This is the continuation of PR which added Clamped Swiglu(used in GPT OSS) support in TE along with Pytorch integration. This PR hooks up the clamped swiglu and dswiglu's nvte APIs to TE Jax.
-
- 29 Sep, 2025 1 commit
-
-
jberchtold-nvidia authored
Address tolerance check for current scaling dact Signed-off-by:
-
- 23 Sep, 2025 1 commit
-
-
Ming-Xu Huang authored
* Adding Amax Primitive and related args. Signed-off-by:
Ming Huang <mingh@nvidia.com> * Enable local-amax for current-scaling and optionally run AR aross FSDP/TP/SP. Signed-off-by:
-
- 15 Sep, 2025 1 commit
-
-
Ming-Xu Huang authored
* Applying the original precision as Norm outputs' and activation compuations. Signed-off-by:
-
- 08 Sep, 2025 1 commit
-
-
Ming-Xu Huang authored
* Fixing few issues with multi-process launching. Signed-off-by:
-
- 05 Sep, 2025 1 commit
-
-
jberchtold-nvidia authored
* Custom call tests passing Signed-off-by:
-
- 03 Sep, 2025 1 commit
-
-
Kshitij Lakhani authored
* Fix failing tests for dropout=0.1 and bias for fused attn for blackwell Signed-off-by:
-
- 27 Aug, 2025 2 commits
-
-
jberchtold-nvidia authored
* Decouple recipe and scaling mode Signed-off-by:
-
Ming-Xu Huang authored
* FP8 AllGather in FP8 GroupedGEMM 1. Support current scaling FP8 quantation with a given amax. 2. Support FP8 AG in fwd and BF16 RS in bwd. 3. The workflow is AR-max -> FP8 Quant -> FP8 AG -> FP8 GroupedGEMM. Signed-off-by:
-
- 26 Aug, 2025 1 commit
-
-
Phuong Nguyen authored
* clean up sharding Signed-off-by:
-
- 25 Aug, 2025 1 commit
-
-
Kshitij Lakhani authored
Add Transformer Layer tests for pre_scale_bias and post_scale_bias Signed-off-by:
-
- 20 Aug, 2025 1 commit
-
-
jberchtold-nvidia authored
[JAX] Error checking for mesh resource and update GemmPrimitive to use global_mesh_resource().fsdp_resource (#2088) * Enforce global MeshResource is set Signed-off-by:
-
- 15 Aug, 2025 1 commit
-
-
Kshitij Lakhani authored
* Move some dist fused attn tests to L2 1. TestReorderCausalLoadBalancing: Run two (non symmetric) BSHD/SBHD data shape combination 2. TestDistributedSelfAttn: Run only one (smaller) BSHD type data shape combination 3. TestDistributedCrossAttn: Run only one (smaller) BSHD type data shape combination 4. TestDistributedContextParallelSelfAttn: Run all cp1 combinations Signed-off-by:
Kshitij Janardan Lakhani <klakhani@login-preos01.a51.clusters.nvidia.com> * Remove test code for L2 dist residing in L2 test.sh Signed-off-by:
-
- 13 Aug, 2025 1 commit
-
-
jberchtold-nvidia authored
* Add L2_jax_distributed_unittest Signed-off-by:
-
- 08 Aug, 2025 1 commit
-
-
Phuong Nguyen authored
* rm cudaGraph compatible trait from GroupedGEMM and groupedQuantize Signed-off-by:
-
- 07 Aug, 2025 1 commit
-
-
Phuong Nguyen authored
* rm batch_dim, sequence_dim, sequence_parallel_output Signed-off-by:
-
- 06 Aug, 2025 1 commit
-
-
jberchtold-nvidia authored
* Pytest timings Signed-off-by:
-
- 24 Jul, 2025 1 commit
-
-
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
-
-
jberchtold-nvidia authored
Fix current scaling test_helper.py and enable test_helper.py in L0 Signed-off-by:
-
- 22 Jul, 2025 1 commit
-
-
Oleg Goncharov authored
* Fixed conflicts 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 * Minor code refactoring to avoid unnecessary checks Signed-off-by:
-
- 21 Jul, 2025 1 commit
-
-
Charlene Yang authored
* exclude 9.10.0/.1 for certain configs Signed-off-by:
Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix kv_channels Signed-off-by:
-
- 19 Jul, 2025 1 commit
-
-
jberchtold-nvidia authored
Update tolerance of distributed layernorm MLP for FP8 Signed-off-by:
-
- 18 Jul, 2025 1 commit
-
-
Phuong Nguyen authored
* set precision=HIGHEST for the ref_grouped_gemm impl in the unit test Signed-off-by:
-
- 17 Jul, 2025 1 commit
-
-
Phuong Nguyen authored
* remove unnecessary padding Signed-off-by:
-
- 15 Jul, 2025 1 commit
-
-
Emmanuel Ferdman authored
* [JAX] Resolve test conflict in JAX helper tests Signed-off-by:
Emmanuel Ferdman <emmanuelferdman@gmail.com> * [JAX] Resolve test conflict in JAX helper tests Signed-off-by:
-
