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

[gemini-code-assist]

Code Review

This pull request introduces support for the GLM-5.2 model, specifically implementing cross-layer IndexShare (IndexCache) for Dynamic Sparse Attention. The changes include new model configurations, layer definitions, and utility functions for managing IndexShare patterns. The reviewer provided critical feedback to improve the robustness of the implementation, specifically suggesting the replacement of hardcoded donor layer logic in checkpoint conversion and attention layers with dynamic lookups, as well as addressing potential indexing errors and log spam.

[gemini-code-assist]

The current implementation hardcodes model.layers.0 as the donor layer for all shared layers. In GLM-5.2, there are multiple Full (F) layers (e.g., 0, 4, 8, 12, ...) acting as donors. Hardcoding layer 0 means higher shared layers (like 5, 6, 7) will incorrectly reuse layer 0's indexer weights instead of their actual donor layer (layer 4), leading to incorrect attention routing and degraded model quality.

We can resolve this by dynamically searching backwards for the closest preceding layer that contains the indexer weights. This is extremely robust and doesn't require access to the config object.

Suggested change

	if "indexer" in str(hf_key_single) and str(hf_key_single).startswith("model.layers."):
	import re
	m = re.match(r"model\.layers\.(\d+)\.(.+)", str(hf_key_single))
	if m:
	rest = m.group(2)
	donor_key = f"model.layers.0.{rest}"
	try:
	hf_tensor_numpy = tensor_getter_fn(donor_key)
	except Exception:
	hf_tensor_numpy = np.zeros(mt_slice_shape, dtype=np.float32)
	else:
	raise e
	else:
	raise e
	if "indexer" in str(hf_key_single) and str(hf_key_single).startswith("model.layers."):
	import re
	m = re.match(r"model\.layers\.(\d+)\.(.+)", str(hf_key_single))
	if m:
	layer_idx = int(m.group(1))
	rest = m.group(2)
	hf_tensor_numpy = None
	for d in range(layer_idx - 1, -1, -1):
	try:
	hf_tensor_numpy = tensor_getter_fn(f"model.layers.{d}.{rest}")
	break
	except Exception:
	continue
	if hf_tensor_numpy is None:
	hf_tensor_numpy = np.zeros(mt_slice_shape, dtype=np.float32)
	else:
	raise e
	else:
	raise e

[gemini-code-assist]

Similar to the issue in _build_single_axis_stacked_tensor, hardcoding model.layers.0 as the donor layer for all shared layers is incorrect for GLM-5.2 because there are multiple Full (F) layers acting as donors. We should use the same backward-search logic here to dynamically find the closest preceding Full layer.

Suggested change

	if "indexer" in key and key.startswith("model.layers."):
	import re
	m = re.match(r"model\.layers\.(\d+)\.(.+)", key)
	if m:
	rest = m.group(2)
	donor_key = f"model.layers.0.{rest}"
	try:
	return orig_tensor_getter(donor_key)
	except Exception:
	pass
	raise e
	if "indexer" in key and key.startswith("model.layers."):
	import re
	m = re.match(r"model\.layers\.(\d+)\.(.+)", key)
	if m:
	layer_idx = int(m.group(1))
	rest = m.group(2)
	for d in range(layer_idx - 1, -1, -1):
	try:
	return orig_tensor_getter(f"model.layers.{d}.{rest}")
	except Exception:
	continue
	raise e

[gemini-code-assist]

Hardcoding is_full = (layer_idx % 4 == 0) assumes the pattern is always "FSSS" and that the layer indexing aligns perfectly. If a different pattern is configured (e.g., "FSS"), this will be incorrect. We should dynamically look up whether the layer is Full using the self.is_full_array we initialized.

Suggested change

	if layer_idx is not None:
	is_full = (layer_idx % 4 == 0)
	indexer_mask, topk_indices, indexer_score = jax.lax.cond(
	is_full,
	_run_full,
	_run_shared,
	operand=None,
	)
	if layer_idx is not None:
	is_full = self.is_full_array[layer_idx]
	indexer_mask, topk_indices, indexer_score = jax.lax.cond(
	is_full,
	_run_full,
	_run_shared,
	operand=None,
	)

[gemini-code-assist]

During scanned execution, self.served_group_size is a static attribute of the template layer, which means it will be constant (e.g., 4) for all layers in the scan. However, the actual group size can vary (e.g., the last group has size 2). We should dynamically look up the group size using self.served_group_size_array to ensure correct loss scaling across all layers.

        if getattr(self.config, "use_index_share", False):
          group_size = self.served_group_size_array[layer_idx] if layer_idx is not None else float(self.served_group_size)
          if group_size > 1:
            loss_scale = loss_scale / group_size

[gemini-code-assist]

Using k.split(".")[2] can raise an IndexError if there are any keys in hf_state_dict_numpy starting with model.layers. but having fewer than 3 components (e.g., model.layers). Using a regular expression match is much safer and more robust.

Suggested change

	matching_layers = [
	int(k.split(".")[2])
	for k in hf_state_dict_numpy
	if k.startswith("model.layers.") and k.endswith(f".{rest}")
	]
	matching_layers = []
	for k in hf_state_dict_numpy:
	m_k = re.match(r"model\.layers\.(\d+)\.(.+)", k)
	if m_k and m_k.group(2) == rest:
	matching_layers.append(int(m_k.group(1)))

[gemini-code-assist]

To support dynamic IndexShare patterns and scanned layers robustly, we should initialize boolean and integer arrays representing whether each layer is Full/Shared and its served group size. This allows us to dynamically index into them using the JAX tracer layer_idx during scanned execution, rather than hardcoding the pattern or assuming static properties.

Suggested change

	is_pruned = (
	getattr(config, "use_index_share", False)
	and getattr(config, "prune_shared_indexers", True)
	and self.is_shared_layer
	)
	if self.use_indexer and not is_pruned:
	is_pruned = (
	getattr(config, "use_index_share", False)
	and getattr(config, "prune_shared_indexers", True)
	and self.is_shared_layer
	)
	if getattr(config, "use_index_share", False):
	from maxtext.utils import index_share_utils
	pattern = index_share_utils.parse_index_share_pattern(config.index_share_pattern, config.num_decoder_layers)
	self.is_full_array = jnp.array([p == "F" for p in pattern], dtype=jnp.bool_)
	self.served_group_size_array = jnp.array(index_share_utils.get_served_group_sizes(pattern), dtype=jnp.float32)
	if self.use_indexer and not is_pruned:

[gemini-code-assist]

Logging on all processes in a multi-host environment can cause significant log spam. We should restrict this log message to process 0.

Suggested change

	if layer_idx == 0:
	num_f = pattern.count("F")
	num_s = pattern.count("S")
	absl.logging.info(
	f"[GLM-5.2 IndexShare Active] Total layers: {config.num_decoder_layers} | "
	f"Pattern: {config.index_share_pattern} | Full (F) layers with active indexers: {num_f} | "
	f"Shared (S) layers with pruned indexers: {num_s} (Pruned {num_s / config.num_decoder_layers * 100:.1f}% indexer compute/parameters)"
	)
	if layer_idx == 0 and jax.process_index() == 0:
	num_f = pattern.count("F")
	num_s = pattern.count("S")
	absl.logging.info(
	f"[GLM-5.2 IndexShare Active] Total layers: {config.num_decoder_layers} | "
	f"Pattern: {config.index_share_pattern} | Full (F) layers with active indexers: {num_f} | "
	f"Shared (S) layers with pruned indexers: {num_s} (Pruned {num_s / config.num_decoder_layers * 100:.1f}% indexer compute/parameters)"
	)