[huytransformer]

Description

This PR adds context_parallel_strategy="ulysses", DeepSpeed-Ulysses style context parallelism for the TPU splash attention (https://arxiv.org/abs/2309.14509).

Motivations: (1) the existing approaches degrade when the context parallelism degree is large relative to the sequence length. (2) It is also the head-parallel half of a Ulysses x ring hybrid (USP), which can be built on top of this PR.

Performance

What we did: v5p-128 (64 chips), bf16, synthetic data, 2M tokens per step in every cell, remat_policy=custom context=device, dq_reduction_steps=3. Each uses its best known block sizes (2048 for Ulysses, 1024 for Ring and All-gather).

llama2-7b (MHA, 32 KV heads), 32K context, cp=16, 2K tokens per device.

Method	Step time	TFLOP/s/device	MFU
Ulysses	8.90 s	239.4	52.2% Xprof
All-gather	8.83 s	242.7	52.9%
Ring	16.38 s	130.8	28.5% Xprof

Ulysses matches all-gather and runs 1.8x faster than ring, while holding 1/16 of all-gather's attention KV per device (2 of 32 heads).

llama3.1-8b (GQA, 8 KV heads), cp=4.

Context (chunk)	Ulysses	All-gather	Ring
64K (16K)	243.3 / 53.0% Xprof	282.5 / 61.6% Xprof	277.4 / 60.4% Xprof
8K (2K)	259.0 / 56.4% Xprof	281.1 / 61.2% Xprof	274.7 / 59.8% Xprof

Correctness: 20 training steps against all_gather on llama3.1-8b at 64K, cp=4, identical data, max loss difference is 0.001.

Example repro

export LIBTPU_INIT_ARGS="--xla_tpu_scoped_vmem_limit_kib=65472 \
  --xla_tpu_use_enhanced_launch_barrier=true \
  --xla_tpu_enable_async_collective_fusion=true \
  --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true \
  --xla_tpu_enable_async_collective_fusion_multiple_steps=true \
  --xla_tpu_overlap_compute_collective_tc=true \
  --xla_enable_async_all_gather=true \
  --xla_enable_async_collective_permute=true \
  --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \
  --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \
  --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true"

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
  base_output_directory=${OUTPUT_DIR} run_name=ulysses_cp16 \
  dataset_type=synthetic enable_checkpointing=false steps=10 \
  model_name=llama2-7b max_target_length=32768 per_device_batch_size=1 \
  ici_context_parallelism=16 ici_fsdp_parallelism=4 \
  attention=flash use_tokamax_splash=true use_jax_splash=false \
  context_parallel_strategy=ulysses context_parallel_load_balance=false \
  packing=false dq_reduction_steps=3 remat_policy=custom context=device \
  num_vocab_tiling=16 \
  sa_block_q=2048 sa_block_kv=2048 sa_block_q_dkv=2048 sa_block_kv_dkv=2048

Tests

• tests/unit/ulysses_attention_test.py: Passed.
• tests/unit/ulysses_collective_test.py: Passed.
• tests/unit/configs_value_test.py: Passed.
• tests/integration/train_tests.py: train smoke test with ulysses cp=4. Passed.

Checklist

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

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• 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, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[github-actions]

🤖 Hi @huytransformer, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This Pull Request introduces highly performant TPU Ulysses context parallelism on the TPU Tokamax Splash attention training path. The implementation is exceptionally clean, well-integrated, and backed by comprehensive unit, collective, and integration tests that verify correctness against dense references and compiled HLO collectives.

🔍 General Feedback

• Excellent Test Coverage: The newly added tests (including CPU multi-device mock checks and HLO collective pattern assertions) set a very high bar for correctness and robustness.
• Robust Config Validation: The comprehensive validation in configs/types.py ensures that unsupported combinations (like MQA, packing, ragged attention, dropout, etc.) are proactively rejected with precise errors.
• Clean Shared Utilities: Factoring out mesh and sharding operations into context_parallel_utils.py is a great design choice that improves maintainability.
• Thorough Documentation: The update to sharding.md is incredibly detailed and accurately captures Ulysses constraints and architectural mechanics.

[codecov]

Codecov Report

❌ Patch coverage is 79.10448% with 42 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_op.py	65.85%	23 Missing and 5 partials ⚠️
src/maxtext/utils/sharding.py	68.00%	5 Missing and 3 partials ⚠️
src/maxtext/kernels/attention/ulysses_attention.py	92.68%	3 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[RissyRan]

Thanks! LGTM at high level. Some minor comments.

[susanbao]

Have added the comments. Please check. It is great to have Ulysses in maxtext.