Code

Anisha seq parallel#1147

Code

Draft

A9isha wants to merge 28 commits intomainAI-Hypercomputer/maxtext:mainfrom anisha-seq-parallelAI-Hypercomputer/maxtext:anisha-seq-parallelCopy head branch name to clipboard

+623-66

Commits

Commits on Jan 22, 2025

• Add sequence parallelism sharding q

  

  raymondzouu

  authored and

  A9isha

  committed
  05ef9c7

  Copy full SHA for 05ef9c7

• Add segmentid implementation

  

  raymondzouu

  authored and

  A9isha

  committed
  c87ecf9

  Copy full SHA for c87ecf9

• add debug statements

  A9isha

  committed
  a965a13

  Copy full SHA for a965a13

• Add segmentid implementation

  

  raymondzouu

  authored and

  A9isha

  committed
  61d5ef8

  Copy full SHA for 61d5ef8

• temp commit to fix splash attention kernel call

  A9isha

  committed
  d01af2e

  Copy full SHA for d01af2e

• fix seq sharding parameter in make_splash_mha

  A9isha

  committed
  a0e093e

  Copy full SHA for a0e093e

• update

  A9isha

  committed
  ce2c42f

  Copy full SHA for ce2c42f

• fix the -1 value for context parallelism

  A9isha

  committed
  ece53a5

  Copy full SHA for ece53a5

• add load_balancing in causal mask

  A9isha

  committed
  54e9a55

  Copy full SHA for 54e9a55

• reorder mask first commit

  A9isha

  committed
  4241570

  Copy full SHA for 4241570

• try to make static

  A9isha

  committed
  567233d

  Copy full SHA for 567233d

• static argnames in progress

  A9isha

  committed
  b70e50b

  Copy full SHA for b70e50b

• made multi_head_mask np.ndarray

  A9isha

  committed
  106ae91

  Copy full SHA for 106ae91

• wrap ndarray

  A9isha

  committed
  442c27a

  Copy full SHA for 442c27a

• fix permuted mask

  A9isha

  committed
  0ca559a

  Copy full SHA for 0ca559a

• fix permutation

  A9isha

  committed
  c283c4f

  Copy full SHA for c283c4f

• clean up

  A9isha

  committed
  4d52390

  Copy full SHA for 4d52390

• fix non load balanced case

  A9isha

  committed
  7feeb07

  Copy full SHA for 7feeb07

• clean up

  A9isha

  committed
  bf2d21a

  Copy full SHA for bf2d21a

• use ComputableMask

  A9isha

  committed
  2edbfcf

  Copy full SHA for 2edbfcf

• fix traced array

  A9isha

  committed
  32539c9

  Copy full SHA for 32539c9

• debugging _ComputableMask but moving to dynamic slice creation

  A9isha

  committed
  b47eb81

  Copy full SHA for b47eb81

• try out new computable mask

  A9isha

  committed
  968f235

  Copy full SHA for 968f235

• try running v6e-256 experiment

  A9isha

  committed
  a026cb1

  Copy full SHA for a026cb1

• clean up config

  A9isha

  committed
  5adc54f

  Copy full SHA for 5adc54f

• try different sa_block values

  A9isha

  committed
  fca5d0e

  Copy full SHA for fca5d0e

Commits on Jan 29, 2025

• codestyle changes

  A9isha

  committed
  b70ea1e

  Copy full SHA for b70ea1e

• revert remaining jnp to np in mask computation

  A9isha

  committed
  a829d23

  Copy full SHA for a829d23