Add support for EP to context parallelism in self-attention

[shuningjin]

Description

Goal

For mixture of expert models, we may use expert parallelism. For attention layer, EP acts as FSDP currently. Built upon previous context parallelism work, this PR is to introduce the option of using EP as CP for attention. This is joint effort with @RissyRan.

FIXES: b/418396648

Main code changes

• attentions.py

  • MHA, MLA, tpu_flash_attention
  • changed logical_axis

• base.yml

  • change logical_axis_rules
  • add a flag to customize expert sharding behavior in attention expert_shard_attention_option: fsdp or context

• unit test: tests.attention_test

  • AttentionTest.test_tpu_flash_attention_cp_and_ep & MLATest.test_tpu_flash_attention_cp_and_ep (extended from cp test)
  • using cp/ep with tpu_flash_attention, with or without context_parallel_load_balance
  • compare logit against dot_product without sharding

Use case

• training, MoE with MHA / MLA attention
• cp_load_balance={true, false}, tpu_flash_attention
• Example (ep_as_cp): ici_expert_parallelism=4, ici_context_parallelism=1, expert_shard_attention_option=context, shard context by 4 in attention, shard expert by 4 for moe
• Example (ep_as_cp + native cp): ici_expert_parallelism=2, ici_context_parallelism=2, expert_shard_attention_option=context, shard context by 4 in attention, shard expert by 2 and context by 2 for moe

Tests

Tested on v5p-8

Verify sharding shape

• end-to-end pretraining: reduced version of Mixtral-8x7b (for MHA) and reduced DeepSeek3-671b (for MLA)
• tpu_flash_kernel, context_parallel_load_balance=True, parallelism: ici_expert_parallelism=2, ici_context_parallelism=2, expert_shard_attention_option=context and ici_expert_parallelism=2, ici_context_parallelism=2, expert_shard_attention_option=fsdp
• See test details in b/418396648#comment7

Verify attention output logit against dot product

• use the added unit test
• tpu_flash_kernel, attention {MHA, MLA}, context_parallel_load_balance={True, False}, parallelism: {ici_expert_parallelism=4, expert_shard_attention_option=context, ici_expert_parallelism=2, expert_shard_attention_option=context, ici_context_parallelism=2}

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.

[RissyRan]

Thanks Shuning! Great work!

[gobbleturk]

Have you considered an approach like conditionally modifying the rules (instead of creating new ones?) This is an approach used for pipeline parallelism

maxtext/MaxText/pyconfig.py

Line 788 in fdf479f

def modify_activation_embed_and_logits_batch(logical_axis_rules):

there are pros and cons of both, both pretty ugly IMO but at least when modifying rules there are

1. less rules
2. no if statements which rules to use

[richjames0]

Really impressive that you understood this and got it working!

[shuningjin]

Resolved merge conflict with nnx migration for attention layer

• Convert Attention and MLA modules to NNX #2066
• Accept inputs shape in Attention/MLA constructor #2089

Re-testing on local v5p-8, diff (before vs. after nnx migration)

• Mini mixtral: 1.1, 1.2. Mini deepseek3: 2.1, 2.2
• Sharding shape is still correct, training loss is close
• TFLOP/s/device is better now. In both cases, the FLOPs calculation has included the recent change. Possible reasons: different docker image (jax 0.6.2 vs. 0.7.0), linen vs. nnx.