Code

[JAX] Sharding Utils#1003

Code

Draft

mingxu1067 wants to merge 3 commits intoNVIDIA:mainNVIDIA/TransformerEngine:mainfrom mingxu1067:mingh/sharding_utilsmingxu1067/TransformerEngine:mingh/sharding_utilsCopy head branch name to clipboard

+136-3

Commits

Commits on Jul 9, 2024

• Adding sharding utils for further development of RingAttention

  Show description for 59a7918
  mingxu1067

  committed
  59a7918

  Copy full SHA for 59a7918

• Fixing E1120 in sharding.py

  Show description for 3df91d5
  mingxu1067

  committed
  3df91d5

  Copy full SHA for 3df91d5

Commits on Jul 15, 2024

• Apply **kwargs to lax_paral_op

  Show description for 5138242
  mingxu1067

  committed
  5138242

  Copy full SHA for 5138242