Onboard glm5.1

[notabee]

Description

Onboards the GLM-5.1 (744B parameters Mixture of Experts) model architecture onto MaxText to support full-parameter pre-training and fine-tuning.

Context & Problem Solved

This change introduces the official GLM-5.1 model configuration parameters to enable pre-training and verification on Google TPU slices. To prevent configuration schema validation errors during execution, the model has been registered as an allowed literal option in the MaxText Pydantic schemas.

Implementation Details

1. Model Configuration: Created src/maxtext/configs/models/glm5.1-744b.yml defining:
   • Multi-head Latent Attention (mla) hyperparameters (q_lora_rank=2048, kv_lora_rank=512, qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256).
   • MoE gating and routing configurations (num_experts=256, num_experts_per_tok=8, shared_experts=1, routed_scaling_factor=2.5).
   • Attention type set to dot_product to satisfy sparse indexer validation.
   • Default activation/weight precision optimized to bfloat16 to reduce HBM memory consumption by 50%.
   • Sanity block layers defaulted to 4 for local single-node verification.
2. Schema Registration: Added "glm5.1-744b" to the allowed ModelName literals in src/maxtext/configs/types.py.

Future Improvements

Scaling to the original full model (78 layers / 744B parameters) requires multi-node TPU slices combined with Pipeline/Tensor/Expert Parallelism sharding configurations.

---

Tests

Environment Setup

If running in a PyTorch-preloaded TPU VM image (e.g., standard vllm-tpu), uninstall conflicting torchvision/triton binary packages to prevent dynamo startup crashes:

python3 -m pip uninstall -y torchvision triton

Reproduction Commands

Verify JAX PJRT device compilation and execute a 2-step synthetic training sanity run (e.g., testing 16 layers, ~131.5B parameters):

export HF_TOKEN="<your_token>"

python3 -m maxtext.trainers.pre_train.train \
    src/maxtext/configs/base.yml \
    model_name=glm5.1-744b \
    base_output_directory=/tmp/maxtext_run \
    run_name=glm5_sanity_run \
    steps=2 \
    per_device_batch_size=1 \
    max_target_length=1024 \
    dataset_type=synthetic \
    tokenizer_type=huggingface \
    tokenizer_path=zai-org/GLM-5.1 \
    hf_access_token=$HF_TOKEN \
    enable_checkpointing=false \
    base_num_decoder_layers=16 \
    override_model_config=true

Verified Scaling Limits (on 8x TPU7x VM, 94.75 GB HBM per device):

• 4 layers (12.98B params): Successful step completion, peak HBM 9.08 GB (9.58% capacity)
• 8 layers (52.49B params): Successful step completion, peak HBM 36.69 GB (38.72% capacity)
• 12 layers (92.00B params): Successful step completion, peak HBM 64.29 GB (67.85% capacity)
• 14 layers (111.75B params): Successful step completion, peak HBM 78.10 GB (82.42% capacity)
• 15 layers (121.63B params): Successful step completion, peak HBM 85.00 GB (89.70% capacity)
• 16 layers (131.51B params): Successful step completion, peak HBM 91.90 GB (96.99% capacity)
• 17 layers (141.39B params): OOM (RESOURCE_EXHAUSTED) during HLO compile (requires 98.79 GB HBM)

---

Checklist

• 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 and provided workload details above.
• I have made or will make corresponding changes to the docs.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!