add glm5.2 indexshare

[notabee]

Description

This PR adds full end-to-end support for GLM-5.2 (744B Mixture of Experts) featuring Cross-Layer IndexShare for Dynamic Sparse Attention (DSA) across training, checkpoint conversion, and inference.

Background & Context

GLM-5.2 is a 744B MoE architecture with 78 layers ($3 \text{ dense} + 75 \text{ MoE}$), 256 routed experts + 1 shared expert (Top-8 routed tokens per token), and Multi-Head Latent Attention (MLA) augmented with a Dynamic Sparse Attention (DSA) Lightning Indexer.

In standard DSA implementations, every transformer layer independently computes indexer query/key projections (wq_b, wk_b) and Top-k indexer scores. However, attention distributions exhibit strong cross-layer similarity across adjacent decoder stages. GLM-5.2 introduces Cross-Layer IndexShare, utilizing an FSSS periodic pattern (1 Full layer followed by 3 Shared layers) where:

• Full ($F$) Layers (20 layers: 0, 4, 8, ...): Actively execute the Lightning Indexer, computing the sparse attention mask and Top-k routing indices.
• Shared ($S$) Layers (58 layers: 1, 2, 3, 5, 6, 7, ...): Prune all indexer parameters and GEMM compute, inheriting and reusing the donor $F$-layer's cached indexer mask.

Key Changes

1. Model Architecture & Configs:
   • Added src/maxtext/configs/models/glm5.2-744b.yml with native support for use_index_share=true, index_share_pattern="FSSS", and prune_shared_indexers=true.
   • Updated src/maxtext/models/glm5.py (GLMDenseLayer, GLMMoELayer, GLMGenericLayer) to thread layer_idx and cached_indexer_state.
2. Scanned Layer Execution & Carry State:
   • Updated _apply_layers_sequentially in src/maxtext/layers/nnx_decoders.py to thread (carry_y, cached_indexer_state, layer_idx) through jax.lax.scan with invariant concrete tensor structures for HBM stability.
   • Updated src/maxtext/layers/attention_mla.py with jax.lax.cond dispatching to conditionally run full indexer calculations on $F$-layers or pass through cached state on $S$-layers.
3. Checkpoint Conversion:
   • Added weight conversion support for HuggingFace zai-org/GLM-5.2 safetensors to MaxText format in to_maxtext.py, automatically pruning 74.4% of indexer weights on shared layers.
4. End-to-End Test Suite:
   • Added tests/end_to_end/tpu/glm5/glm5.2-744b/1_test_glm5.sh (Checkpoint Conversion).
   • Added tests/end_to_end/tpu/glm5/glm5.2-744b/2_test_glm5.sh (Pre-Training & Generation).

Benefits

• Compute & Parameter Efficiency: Prunes 74.4% of indexer parameters and FLOPs across the model (saving 58 indexer projections per token).
• HBM Footprint: Fits cleanly within TPU v7x HBM (~27.5 GB / 94.74 GB during scanned pre-training).

---

Tests

1. Checkpoint Conversion (Step 1)

bash tests/end_to_end/tpu/glm5/glm5.2-744b/1_test_glm5.sh

2. Distributed Pre-Training & Verification on TPU v7x (32 Chips / 64 Devices)

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
    base_output_directory="gs://maxtext-glm5-europe-west4/GLM-5.2/training_logs" \
    run_name=glm52_pretrain_verify \
    model_name=glm5.2-744b \
    scan_layers=true \
    use_indexer=true \
    use_index_share=true \
    index_share_pattern="FSSS" \
    prune_shared_indexers=true \
    indexer_sparse_training=true \
    tokenizer_type=huggingface \
    tokenizer_path=zai-org/GLM-5.2 \
    dataset_type=synthetic \
    per_device_batch_size=1 \
    max_target_length=4096 \
    ici_expert_parallelism=4 \
    ici_fsdp_parallelism=16 \
    steps=10

3. Autoregressive Generation & Decoding

python3 -m maxtext.inference.decode src/maxtext/configs/base.yml \
    model_name=glm5.2-744b \
    tokenizer_type=huggingface \
    tokenizer_path=zai-org/GLM-5.2 \
    load_parameters_path="gs://maxtext-glm5-europe-west4/maxtext-glm-5.2-bf16-converted-final-78l/0/items" \
    scan_layers=true \
    use_indexer=true \
    use_index_share=true \
    index_share_pattern="FSSS" \
    prune_shared_indexers=true \
    prompt="The capital of France is"

---

Checklist

• 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 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).