Skip to content

Add TPU USP context parallelism #4836

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
huytransformer wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add TPU USP context parallelism #4836

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/guides/optimization/sharding.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -262,6 +262,8 @@ MaxText supports `context_parallel_strategy=all_gather`, and supports `context_p

MaxText also supports `context_parallel_strategy=ulysses` ([DeepSpeed Ulysses](https://arxiv.org/abs/2309.14509)) on the TPU Tokamax Splash path for training. Ulysses exchanges sequence ownership for head ownership by communicating the Q, K, V, and output activations through all-to-all collectives: each device computes ordinary full-sequence attention for its head subset, and the inverse all-to-all restores the sequence sharding on the output. It requires explicit positive context parallelism values, `context_sharding=context`, `attention=flash` with Tokamax Splash, global causal attention, query and KV head counts divisible by the context parallel size including after tensor-parallel head sharding, matching Q and KV head-sharding axes, an unsharded head feature dimension, a divisible sequence length, `dq_reduction_steps` of 0 or 3, `context_parallel_load_balance=false` (each device computes full-sequence attention for its head subset, so the work is already balanced and the causal load-balancing reorder must stay off), and ICI-only context parallelism (`dcn_context_parallelism` must equal 1). It does not support MQA, dropout, QK-Clip statistics, ragged attention, attention sinks, sparse indexer masks, chunked prefill, MoBA, or multimodal attention.

MaxText also supports `context_parallel_strategy=usp` ([USP](https://arxiv.org/abs/2405.07719), Ulysses over ring) on the same TPU Tokamax Splash path for training. This initial support is non-load-balanced. USP factors the context parallelism into a ring dimension on the `context` mesh axis (`ici_context_parallelism`) and a Ulysses dimension on the `context_ulysses` mesh axis (`ici_context_ulysses_parallelism`): the Ulysses all-to-all exchanges sequence ownership for head ownership over the Ulysses axis at each fixed ring position, and the ring kernel then rotates K and V across the ring axis inside each head subset. The strategy is hybrid-only: both dimensions must be greater than one, and the single-dimension endpoints are the existing `ring` and `ulysses` strategies. It shares the Ulysses restrictions (explicit positive ICI-only sizes, `attention=flash` with Tokamax Splash, global causal attention, head counts divisible by the Ulysses size, no MQA, no packing, no load balancing, no dropout, no multi-token prediction, no dKV megacore) and additionally requires `max_target_length` divisible by the total context parallelism and by the ring size squared.

### CP Arithmetic Intensity

The main communications are the same as FSDP (all gather weights and synchronize gradients), with an arithmetic intensity of `local_batch` / `sparsity`.
Expand Down
47 changes: 26 additions & 21 deletions src/maxtext/configs/base.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -514,30 +514,30 @@ compile_xla_flags: "" # Compiler options e.g. compile_xla_flags="--xla_tpu_num_s
# Parallelism
shard_mode: "auto" # can be either auto or explicit
custom_mesh_and_rule: "" # replace default mesh and logical rule by specifying yml name under config/mesh_and_rule/.
mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_sequence', 'expert', 'autoregressive']
mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_ulysses', 'context_autoregressive', 'tensor', 'tensor_sequence', 'expert', 'autoregressive']
logical_axis_rules: [
['circular_repeats', []],
# ==========================================
# Vocabulary Embedding
# ==========================================
# Vocab Activations
['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert']],
['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'expert']],
['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_ulysses', 'expert']],
['activation_vocab', ['tensor', 'tensor_sequence']],
['activation_vocab', ['tensor']],
['activation_vocab', 'tensor_sequence'],
# Vocab Weights
['vocab', ['tensor', 'tensor_sequence', 'autoregressive']],
['embed_vocab', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
['embed_vocab', ['fsdp', 'fsdp_transpose', 'context', 'context_ulysses', 'expert']],
# ==========================================
# Attention
# ==========================================
# Attention Activations
['activation_batch_attn', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
['activation_heads', ['tensor', 'tensor_sequence', 'autoregressive']],
['activation_kv_heads', ['tensor', 'tensor_sequence']],
['activation_length_attn', ['context']],
['activation_q_length', ['context']],
['activation_length_attn', ['context', 'context_ulysses']],
['activation_q_length', ['context', 'context_ulysses']],
['activation_kv_length', []],
['activation_embed_attn', ['tensor']],
['activation_kv', ['tensor', 'tensor_sequence']],
Expand All @@ -550,45 +550,45 @@ logical_axis_rules: [
['qkv', []],
['kv', []],
['kv_head_dim', []],
['q_lora', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
['q_lora', ['fsdp', 'context', 'expert']],
['q_lora', ['fsdp', 'fsdp_transpose', 'context', 'context_ulysses', 'expert']],
['q_lora', ['fsdp', 'context', 'context_ulysses', 'expert']],
["q_lora_up_proj", []],
['kv_lora', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
['kv_lora', ['fsdp', 'context', 'expert']],
['kv_lora', ['fsdp', 'fsdp_transpose', 'context', 'context_ulysses', 'expert']],
['kv_lora', ['fsdp', 'context', 'context_ulysses', 'expert']],
["kv_lora_up_proj", []],
# ==========================================
# Mixture of Experts (MoE)
# ==========================================
# MoE Activations
['activation_batch_moe', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
['activation_length_moe', ['context']],
['activation_norm_length_moe', ['tensor_sequence', 'context']],
['activation_length_moe', ['context', 'context_ulysses']],
['activation_norm_length_moe', ['tensor_sequence', 'context', 'context_ulysses']],
['activation_embed_moe', ['tensor']],
['activation_mlp_moe', ['tensor', 'tensor_sequence']],
['activation_exp', ['expert']],
# MoE Weights
['exp', 'expert'],
['mlp_moe', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
['embed_moe', ['fsdp', 'fsdp_transpose', 'context']],
['embed_moe', ['fsdp', 'context']],
['embed_moe', ['fsdp', 'fsdp_transpose', 'context', 'context_ulysses']],
['embed_moe', ['fsdp', 'context', 'context_ulysses']],
# ==========================================
# Standard MLP / Dense Layers / Model Structure
# ==========================================
# Dense Activations
['activation_mlp', ['tensor', 'tensor_sequence']],
# Note activation batch and length also get used in vocab
['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
['activation_length', ['context']],
['activation_norm_length', ['tensor_sequence', 'context']],
['activation_length', ['context', 'context_ulysses']],
['activation_norm_length', ['tensor_sequence', 'context', 'context_ulysses']],
['activation_embed', ['tensor']],
['activation_stage', 'stage'],
# General Weights
['mlp', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
# GDN (linear-attention) projections shard like 'mlp' during training; the
# vLLM serving config overrides this to match tpu-inference's ATTN_HEAD order.
['gdn_head', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
['embed', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
['embed', ['fsdp', 'context', 'expert']],
['embed', ['fsdp', 'fsdp_transpose', 'context', 'context_ulysses', 'expert']],
['embed', ['fsdp', 'context', 'context_ulysses', 'expert']],
['norm', ['tensor']],
['layers', 'stage'],
['diloco', 'diloco'],
Expand All @@ -600,8 +600,8 @@ logical_axis_rules: [
# ==========================================
# Inference(Prefill, Decode, Cache)
# ==========================================
['prefill_activation_length', ['context']],
['prefill_activation_norm_length', ['tensor_sequence', 'context']],
['prefill_activation_length', ['context', 'context_ulysses']],
['prefill_activation_norm_length', ['tensor_sequence', 'context', 'context_ulysses']],
['activation_prefill_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
['decode_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
['decode_length', []],
Expand All @@ -622,11 +622,14 @@ logical_axis_rules: [
['exp_with_fsdp', 'fsdp'],
]
# Axes used for DCN must be earlier in this list than ICI, see (b/339009148) for details
data_sharding: [['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_sequence', 'expert', 'autoregressive']]
data_sharding: [['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_ulysses', 'context_autoregressive', 'tensor', 'tensor_sequence', 'expert', 'autoregressive']]
input_data_sharding_logical_axes: ['activation_embed_and_logits_batch', 'activation_norm_length']
# Determines which physical axis plays the role of context parallelism for input data processing and load balancing
# only supports "context" or "expert" (when custom_mesh_and_rule=ep-as-cp)
context_sharding: "context"
# Physical axis for the Ulysses (all-to-all) dimension of context_parallel_strategy='usp'
# only supports "context_ulysses"
ulysses_context_sharding: "context_ulysses"
# Customized mesh and logical rules for evaluation (e.g. "ep-as-cp")
custom_mesh_and_rule_for_eval: ""
logical_axis_rules_for_eval: []
Expand All @@ -644,6 +647,7 @@ dcn_fsdp_parallelism: 1
dcn_fsdp_transpose_parallelism: 1
dcn_sequence_parallelism: 1 # never recommended
dcn_context_parallelism: 1
dcn_context_ulysses_parallelism: 1
dcn_context_autoregressive_parallelism: 1
dcn_tensor_parallelism: 1 # never recommended
dcn_tensor_sequence_parallelism: 1 # never recommended
Expand All @@ -656,6 +660,7 @@ ici_fsdp_parallelism: -1 # recommended ICI axis to be auto-sharded
ici_fsdp_transpose_parallelism: 1
ici_sequence_parallelism: 1
ici_context_parallelism: 1
ici_context_ulysses_parallelism: 1 # Ulysses (head-exchange) dimension of the context parallelism; used by context_parallel_strategy='usp'.
ici_context_autoregressive_parallelism: 1
ici_tensor_parallelism: 1
ici_tensor_sequence_parallelism: 1
Expand Down Expand Up @@ -1154,7 +1159,7 @@ cost_estimate_flops_bwd: -1 # -1 means using splash default cost estmiation, any
dq_reduction_steps: 0 #the number of reduction steps. For now, only 3 or all the kv steps are supported.
### Determine if we want to use load balance for context parallelism
context_parallel_load_balance: true
context_parallel_strategy: "all_gather" # "all_gather", "ring", or "ulysses"
context_parallel_strategy: "all_gather" # "all_gather", "ring", "ulysses", or "usp"
context_parallel_reorder_strategy: "auto" # "auto", "dual_chunk_swap", or "striped"


Expand Down
Loading
Loading