Code

Anisha seq parallel#1147

Code

Draft

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

+479-64

Commits

Commits on Dec 15, 2024

• Add sequence parallelism sharding q

  raymondzouu

  committed
  7e81fec

  Copy full SHA for 7e81fec

Commits on Dec 17, 2024

• Add segmentid implementation

  raymondzouu

  committed
  e46a5ad

  Copy full SHA for e46a5ad

Commits on Jan 3, 2025

• Add segmentid implementation

  raymondzouu

  committed
  6f6176d

  Copy full SHA for 6f6176d

• add debug statements

  A9isha

  committed
  c790881

  Copy full SHA for c790881

• Merge branch 'raymondzou-context-parallelism' of https://github.com/google/maxtext into raymondzou-context-parallelism

  A9isha

  committed
  60873ee

  Copy full SHA for 60873ee

Commits on Jan 4, 2025

• temp commit to fix splash attention kernel call

  A9isha

  committed
  18c691b

  Copy full SHA for 18c691b

Commits on Jan 6, 2025

• fix seq sharding parameter in make_splash_mha

  A9isha

  committed
  002eaea

  Copy full SHA for 002eaea

Commits on Jan 7, 2025

• update

  A9isha

  committed
  1023c94

  Copy full SHA for 1023c94

• fix the -1 value for context parallelism

  A9isha

  committed
  a4dd7b7

  Copy full SHA for a4dd7b7

Commits on Jan 8, 2025

• add load_balancing in causal mask

  A9isha

  committed
  1168042

  Copy full SHA for 1168042

• reorder mask first commit

  A9isha

  committed
  a7c2e5b

  Copy full SHA for a7c2e5b

Commits on Jan 9, 2025

• try to make static

  A9isha

  committed
  0f07893

  Copy full SHA for 0f07893

• static argnames in progress

  A9isha

  committed
  33da5db

  Copy full SHA for 33da5db

• made multi_head_mask np.ndarray

  A9isha

  committed
  92230da

  Copy full SHA for 92230da

Commits on Jan 10, 2025

• wrap ndarray

  A9isha

  committed
  e75c257

  Copy full SHA for e75c257

• fix permuted mask

  A9isha

  committed
  bdd5b5c

  Copy full SHA for bdd5b5c

• fix permutation

  A9isha

  committed
  d3dc0f9

  Copy full SHA for d3dc0f9

• clean up

  A9isha

  committed
  13f9aab

  Copy full SHA for 13f9aab

• fix non load balanced case

  A9isha

  committed
  c9ab25c

  Copy full SHA for c9ab25c

• clean up

  A9isha

  committed
  29b69ad

  Copy full SHA for 29b69ad

• use ComputableMask

  A9isha

  committed
  a87f9db

  Copy full SHA for a87f9db

Commits on Jan 11, 2025

• fix traced array

  A9isha

  committed
  be2300f

  Copy full SHA for be2300f

Commits on Jan 13, 2025

• debugging _ComputableMask but moving to dynamic slice creation

  A9isha

  committed
  dcdbb67

  Copy full SHA for dcdbb67