Code

[pull] main from NVIDIA:main#26

Code

Merged

phu0ngng merged 63 commits intophu0ngng:mainphu0ngng/TransformerEngine:mainfrom NVIDIA:mainNVIDIA/TransformerEngine:mainCopy head branch name to clipboardOct 25, 2024

+6,620-2,128

Commits

Commits on Sep 16, 2024

• [JAX] Fix unit tests to work around cuDNN 9.4 regression of 0 length sequences (#1179)

  Show description for df69965
  mgoldfarb-nvidia

  authored
  df69965

  Copy full SHA for df69965

• Add dtensor support for TE optimizers (#1171)

  Show description for af5daa0
  

  blahBlahhhJ

  and

  ksivaman

  authored
  af5daa0

  Copy full SHA for af5daa0

• Update CI users (#1181)

  Show description for d2d4cf9
  timmoon10

  authored
  d2d4cf9

  Copy full SHA for d2d4cf9

Commits on Sep 17, 2024

• [JAX] Context Parallel Attention with All-Gather (#1106)

  Show description for 9101a78
  mgoldfarb-nvidia

  authored
  9101a78

  Copy full SHA for 9101a78

• [Common] Default CUDA_HOME to /usr/local/cuda when dynamically loading cuDNN and NVRTC (#1183)

  Show description for 44fd316
  denera

  authored
  44fd316

  Copy full SHA for 44fd316

• Changed VERSION to 1.12.0.dev

  Show description for 528d44b
  ptrendx

  committed
  528d44b

  Copy full SHA for 528d44b

• Allow specifying cmake setup directory (#1186)

  Show description for 28f95bd
  

  ryxli

  and

  Ryan Li

  authored
  28f95bd

  Copy full SHA for 28f95bd

• Add docs for installing from PyPI (#1184)

  Show description for eb60b1a
  ksivaman

  authored
  eb60b1a

  Copy full SHA for eb60b1a

Commits on Sep 18, 2024

• [PyTorch] Port fused optimizer tests to pytest (#1185)

  Show description for 7e1068b
  timmoon10

  authored
  7e1068b

  Copy full SHA for 7e1068b

• [PyTorch] Check network interface name when initializing Userbuffers (#1175)

  Show description for 841634c
  

  3 peopleauthored
  841634c

  Copy full SHA for 841634c

• Expose rotary_base as an arg instead of hardcoding (#944)

  Show description for c0caadb
  

  sudhakarsingh27

  and

  timmoon10

  authored
  c0caadb

  Copy full SHA for c0caadb

Commits on Sep 19, 2024

• [PyTorch] Relax the contiguous check for flash attention (#1176)

  Show description for 0ee5ccd
  yaox12

  authored
  0ee5ccd

  Copy full SHA for 0ee5ccd

Commits on Sep 20, 2024

• Allow downloading of model weights automatically (#1172)

  Show description for 195d703
  

  sudhakarsingh27

  and

  pre-commit-ci[bot]

  authored
  195d703

  Copy full SHA for 195d703

• Restore compatibility with Python 3.8 (#1189)

  Show description for 0c74535
  

  ptrendx

  and

  pre-commit-ci[bot]

  authored
  0c74535

  Copy full SHA for 0c74535

Commits on Sep 23, 2024

• Update list of CI users (#1198)

  Show description for a68acd7
  timmoon10

  authored
  a68acd7

  Copy full SHA for a68acd7

Commits on Sep 24, 2024

• Allow to pass architectures like 90a, without being overriden (#1178)

  Show description for 99af5c0
  

  aurianer

  and

  timmoon10

  authored
  99af5c0

  Copy full SHA for 99af5c0

• Update list of CI users (#1203)

  Show description for a44cb72
  ksivaman

  authored
  a44cb72

  Copy full SHA for a44cb72

Commits on Sep 25, 2024

• fix NVTE_UB_WITH_MPI read (#1194)

  Show description for 209b8e5
  

  erhoo82

  and

  ksivaman

  authored
  209b8e5

  Copy full SHA for 209b8e5

Commits on Sep 27, 2024

• [PyTorch] Add GroupedLinear to the docs and fix typos (#1206)

  Show description for c4a5cb8
  pggPL

  authored
  c4a5cb8

  Copy full SHA for c4a5cb8

• [PyTorch] Fix detection of 3 in 3hd/h3d layouts (#1187)

  Show description for 8a1b7ee
  

  cyanguwa

  and

  pre-commit-ci[bot]

  authored
  8a1b7ee

  Copy full SHA for 8a1b7ee

• Fix CP unit test on A100 and L40s (#1211)

  Show description for 7b152a8
  xrennvidia

  authored
  7b152a8

  Copy full SHA for 7b152a8

Commits on Oct 1, 2024

• [PyTorch] Add pool argument to make_graphed_callable (#1218)

  Show description for 728c558
  ksivaman

  authored
  728c558

  Copy full SHA for 728c558

• [PyTorch] Fix distributed testing (#1219)

  Show description for 46075b9
  ksivaman

  authored
  46075b9

  Copy full SHA for 46075b9

• Removed the unused options from GroupedLinear docs and fixed the bug with offsets (#1220)

  Show description for fb74961
  

  ptrendx

  and

  pre-commit-ci[bot]

  authored
  fb74961

  Copy full SHA for fb74961

Commits on Oct 3, 2024

• [PyTorch] Move block_table argument to FA varlen function (#1222)

  Show description for 10cceae
  cyanguwa

  authored
  10cceae

  Copy full SHA for 10cceae

Commits on Oct 4, 2024

• [PyTorch] Minor optimizations to reduce CPU overheads in modules (#1191)

  Show description for 9d976bc
  

  timmoon10

  and

  pre-commit-ci[bot]

  authored
  9d976bc

  Copy full SHA for 9d976bc

Commits on Oct 6, 2024

• [PyTorch] remove duplicate code (#1215)

  Show description for f8eb799
  

  emmanuel-ferdman

  and

  ksivaman

  authored
  f8eb799

  Copy full SHA for f8eb799

Commits on Oct 7, 2024

• Tests for distributed (#1196)

  Show description for 60f738f
  

  4 peopleauthored
  60f738f

  Copy full SHA for 60f738f

• Hierarchical CP implementation (Ulysses + Ring) (#1209)

  Show description for c24a4c4
  

  xrennvidia

  and

  pre-commit-ci[bot]

  authored
  c24a4c4

  Copy full SHA for c24a4c4

• Fix cuDNN sliding window size (#1212)

  Show description for c3b3cd2
  

  cyanguwa

  and

  ksivaman

  authored
  c3b3cd2

  Copy full SHA for c3b3cd2

Commits on Oct 8, 2024

• [PyTorch] Miscellaneous fixes for FA3 attention (#1174)

  Show description for e762592
  

  cyanguwa

  and

  pre-commit-ci[bot]

  authored
  e762592

  Copy full SHA for e762592

Commits on Oct 9, 2024

• [PyTorch] Debug dtype casting in operation-based API (#1202)

  Show description for 5b89f1a
  

  timmoon10

  and

  pre-commit-ci[bot]

  authored
  5b89f1a

  Copy full SHA for 5b89f1a

• [PyTorch] Add documentation for FP8 attention checkpointing (#1223)

  Show description for 2d87552
  

  cyanguwa

  and

  pre-commit-ci[bot]

  authored
  2d87552

  Copy full SHA for 2d87552

• [PyTorch] Improve get_qkv_layout (#1214)

  Show description for 5b6546c
  

  cyanguwa

  and

  pre-commit-ci[bot]

  authored
  5b6546c

  Copy full SHA for 5b6546c

Commits on Oct 10, 2024

• [JAX] Expose sliding window attn to TE-JAX API (#1205)

  Show description for 85e60e6
  

  3 peopleauthored
  85e60e6

  Copy full SHA for 85e60e6

• Small fixes to Float8Tensor (#1225)

  Show description for 3b89c36
  

  ptrendx

  and

  pre-commit-ci[bot]

  authored
  3b89c36

  Copy full SHA for 3b89c36

Commits on Oct 11, 2024

• Fix bug in torch compile and seqdim is integer (#1217)

  Show description for 9ee2dbd
  

  wplf

  and

  ksivaman

  authored
  9ee2dbd

  Copy full SHA for 9ee2dbd

• Add FlashAttention3 to CP implementations (#1232)

  Show description for b36bd0a
  

  3 peopleauthored
  b36bd0a

  Copy full SHA for b36bd0a

Commits on Oct 12, 2024

• [PyTorch] Let Fused RoPE support CP with THD format (#1238)

  Show description for 55dcbb4
  

  yaox12

  and

  xrennvidia

  authored
  55dcbb4

  Copy full SHA for 55dcbb4

Commits on Oct 14, 2024

• Do not link against CUDA driver when building (#1240)

  Show description for 86f07be
  timmoon10

  authored
  86f07be

  Copy full SHA for 86f07be

Commits on Oct 15, 2024

• Check for backend support in Jax context parallel fused attention test (#1227)

  Show description for 20c55e4
  mgoldfarb-nvidia

  authored
  20c55e4

  Copy full SHA for 20c55e4

• Create README.md for examples/ (#1221)

  Show description for 54aa12a
  

  sbhavani

  and

  ksivaman

  authored
  54aa12a

  Copy full SHA for 54aa12a

Commits on Oct 16, 2024

• [PyTorch] Build custom ORT ops before running ONNX export tests (#1252)

  Show description for f6b766b
  timmoon10

  authored
  f6b766b

  Copy full SHA for f6b766b

• fix assertion bug for SWA API in TE-JAX (#1242)

  Show description for 43b9e1e
  

  3 peopleauthored
  43b9e1e

  Copy full SHA for 43b9e1e

• [PyTorch] Drop FA as an installation requirement (#1226)

  Show description for 161b1d9
  

  cyanguwa

  and

  pre-commit-ci[bot]

  authored
  161b1d9

  Copy full SHA for 161b1d9

• Upgrade pylint to 3.3.1 (#1257)

  Show description for 6e90fcb
  ksivaman

  authored
  6e90fcb

  Copy full SHA for 6e90fcb

• [PyTorch] Fix FP8 activation recompute (#1254)

  Show description for a518151
  ksivaman

  authored
  a518151

  Copy full SHA for a518151

• Changed VERSION to 1.13.0.dev

  Show description for 9001081
  ptrendx

  committed
  9001081

  Copy full SHA for 9001081

Commits on Oct 17, 2024

• [PyTorch] Fix wgrads for GroupedLinear when weights don't require grad (#1258)

  Show description for 2d7020e
  

  yaox12

  and

  ksivaman

  authored
  2d7020e

  Copy full SHA for 2d7020e

• [Bugfix] Fix bias for 0-dim tensors in gemm (#1246)

  Show description for 8e97c8d
  yaox12

  authored
  8e97c8d

  Copy full SHA for 8e97c8d

• [TE/JAX] Enabling CudaGraph for custom calls with FFI (#1228)

  Show description for 12f30ea
  phu0ngng

  authored
  12f30ea

  Copy full SHA for 12f30ea

• Fix seq_dim in CP implementation (#1264)

  Show description for a488b8b
  xrennvidia

  authored
  a488b8b

  Copy full SHA for a488b8b

Commits on Oct 18, 2024

• [PyTorch] Reorganize L1 tests (#1255)

  Show description for 41fe1e5
  timmoon10

  authored
  41fe1e5

  Copy full SHA for 41fe1e5

• [Paddle] Debug wheel test (#1265)

  Show description for 927bca7
  timmoon10

  authored
  927bca7

  Copy full SHA for 927bca7

• [PyTorch] Remove PyTorch L0 distributed test (#1273)

  Show description for 3ea7dd3
  timmoon10

  authored
  3ea7dd3

  Copy full SHA for 3ea7dd3

Commits on Oct 21, 2024

• [PyTorch] Reduce the number of FA versions in L3 tests (#1280)

  Show description for 29e3a09
  cyanguwa

  authored
  29e3a09

  Copy full SHA for 29e3a09

Commits on Oct 22, 2024

• Fused Attention Support 64-bit Ragged Offsets for Large THD Tensors (#1230)

  Show description for 7b18f23
  mgoldfarb-nvidia

  authored
  7b18f23

  Copy full SHA for 7b18f23

• [JAX] Skip V100 encoder tests (#1262)

  Show description for 35f7d26
  zlsh80826

  authored
  35f7d26

  Copy full SHA for 35f7d26

• Add THD + GQA supports (#1260)

  Show description for d9b4bfb
  zlsh80826

  authored
  d9b4bfb

  Copy full SHA for d9b4bfb

Commits on Oct 24, 2024

• [JAX] Fix correctness of JAX fused attention with CP and improve numerics check in unit tests (#1282)

  Show description for 20c7529
  mgoldfarb-nvidia

  authored
  20c7529

  Copy full SHA for 20c7529

• [JAX] XLA Custom Calls with FFI for FusedAttnFwd, Quantize, Transpose, ActLuFP8, LayerNormForwardFP8FFI, and LayerNormBackwardFFI (#1263)

  Show description for 18c2234
  

  huanghua1994

  and

  pre-commit-ci[bot]

  authored
  18c2234

  Copy full SHA for 18c2234

• [Paddle] Update type names for Paddle 3.0 (#1286)

  Show description for 7a5fd0c
  timmoon10

  authored
  7a5fd0c

  Copy full SHA for 7a5fd0c

• [Pytorch] Check gradient in test numerics (#1229)

  Show description for 7b284fe
  

  4 peopleauthored
  7b284fe

  Copy full SHA for 7b284fe