[JAX] Support SP + RoPE + GeLU - #602

Merged

denera merged 20 commits into

NVIDIA:mainfrom

mingxu1067:mingh/sp_rope_gelu

Feb 2, 2024

Commits on Jan 14, 2024

Adding support of sequence parallelism
mingxu1067
committed
3a404ca
Adding RoPE
mingxu1067
committed
604dcf2
Fix wrong batch_logical_axes
mingxu1067
committed
426050d
Rnaming FSDP outer env var
mingxu1067
committed
6f6a927
Poring RoPE to Praxis layers.
mingxu1067
committed
092205f
Porting GeLU + [FP8 Cast].
mingxu1067
committed
66a9b9c
WAR to make XLA successfully match FP8 GEMM on FFN1 with GeLU.
mingxu1067
committed
0e5d703
Allowing arbitrary dimension of NVShape for the workspace allocation
mingxu1067
committed
6693bad
Adding checkpoint_name to fused functions of mlp.py to get better perf with nn.scan.
mingxu1067
committed
90793a9

Commits on Jan 16, 2024

Modify with review feedback.
mingxu1067
committed
e7a447a

Commits on Jan 17, 2024

Merge remote-tracking branch 'upstream/main' into mingh/sp_rope_gelu
mingxu1067
committed
210a8c0

Commits on Jan 18, 2024

Fix bugs
mingxu1067
committed
3b2e231

Commits on Jan 19, 2024

Fix typo.
mingxu1067
committed
780a5f6
Fixed for lint
mingxu1067
committed
6bcdaf6

Commits on Jan 22, 2024

Follow review feedback to modify code.
mingxu1067
committed
6414b29

Commits on Jan 25, 2024

Fix typo.
mingxu1067
committed
f685393

Commits on Jan 30, 2024

Merge remote-tracking branch 'upstream/main' into mingh/sp_rope_gelu
mingxu1067
committed
4222c64

Commits on Jan 31, 2024

Port SP to Praxis
mingxu1067
committed
2ea411a
Fix an issue when enabling both GQA and RoPE.
mingxu1067
committed
45993f6
Update docs
mingxu1067
committed
8e7dcdf