Add context parallelism

[A9isha]

Description

Adding Context Parallelism to MaxText to support long context length in training on TPUs, following PR 1133 which adds it for GPUs
This is because we want to support the common context parallelism paradigm of AG keys and values, shard on query sequence.
Previously, MaxText implemented sequence parallelism by sharding the attention operations on heads (the FF operation being sharded by sequence). This decision was made because the heads dim acts like a batch dimension. This PR's solution frees heads dim to be used for other model parallelism.

We also ensure that we employ "load balancing" in the context parallelism - meaning, we reorder the tokens and the attention mask such that in attention all devices have somewhat similar amount of work (and hence time taken) other wise, with say context parallelism=2, device which receives shard with say token 0,1,2,3 have much lesser work than device which received shard with say token 4,5,6,7, so we load balance by dividing the token as 0,1,6,7 and 2,3,4,5

We use jnp.take in the load balancing, which hopefully can be improved and is being tracked in b/413770626

FIXES: b/377904983

Tests

Tested locally on v5p-8,

and also on v6e-256 for Llama3.1-70b with context length of 131072 tokens

1. using python3 -m benchmarks.maxtext_xpk_runner
2. by updating the model to be used as llama3_1_70b_131072 from maxtext/benchmarks/maxtext_trillium_model_configs.py

Also added unit test for context parallelism in MaxText/tests/attention_test.py

Updated maxtext/end_to_end/tpu/llama2/7b/test_llama2_7b.sh to double as integration test for context parallelism for training in train.py via ici_context_parallelism.
Also, used ici_context_parallelism for forward_pass_logit_checker to test forward pass since we can't use flash attention in decode.py

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[gobbleturk]

go anisha go

[gobbleturk]

Approved! Please don't change the _DENSE_VMEM_LIMIT in this PR though, instead we can document using a different VMEM value in our MaxText configs

[richjames0]

lgtm

[richjames0]

thanks for the change!

[richjames0]

lgtm again