Onboard GLM-5.2 - #4972
Conversation
…p gcloud call in try-except
…eaved DeepSeekV4 RoPE broadcasting, and C++ symbol collision prevention
…c_tpu_v2_kernel_test to fix CPU unit test collection
…_kernel_test to satisfy pyink
…__init__ that broke non-torch CI test collection
… preserve DeepSeekV4RotaryEmbedding in embeddings
…b projection hooks
# Conflicts: # src/maxtext/layers/nnx_decoders.py # tests/unit/nnx_decoders_test.py
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for the GLM-5.1 and GLM-5.2 models, including Mixture of Experts (MoE) and a cross-layer IndexShare (IndexCache) mechanism. Key changes include new GLM layer definitions, updates to the MLA attention layer to support index sharing, and corresponding configuration and test additions. The review feedback highlights several critical issues: a potential JAX tracing error when cached_indexer_state is None in jax.lax.cond, a missing initialization of self.indexer leading to an AttributeError on pruned layers, and a PyTree structural mismatch when scanning layers with pruned indexers. Additionally, suggestions were made to avoid regressions by keeping the default unsqueeze_dim in RotaryEmbedding, pre-computing JAX constant arrays in __init__, correcting the profiler's active state checks for periodic profiling, and fixing slot slicing in the gcloud_stub tokenizer.
| def _run_shared(_): | ||
| with jax.named_scope("glm_shared_layer_index_reuse"): | ||
| mask, indices, score = cached_indexer_state | ||
| mask = checkpoint_name(mask, "shared_layer_reused_mask") | ||
| indices = checkpoint_name(indices, "shared_layer_reused_indices") | ||
| return mask, indices, score |
There was a problem hiding this comment.
During JAX tracing of the jax.lax.cond block, both the _run_full and _run_shared branches are traced. If cached_indexer_state is None (which is the case for the first layer or during initial tracing), unpacking it in _run_shared will raise a TypeError: cannot unpack non-iterable NoneType object. To prevent this tracing error, add a safe fallback to return dummy zero tensors of the correct shape and dtype when cached_indexer_state is None.
def _run_shared(_):
with jax.named_scope("glm_shared_layer_index_reuse"):
if cached_indexer_state is None:
batch = query.shape[0]
q_len = query.shape[1]
kv_len = key.shape[1] if key is not None else 0
topk = self.config.indexer_topk
mask = jnp.zeros((batch, q_len, kv_len), dtype=jnp.bfloat16)
indices = jnp.zeros((batch, q_len, topk), dtype=jnp.int32)
score = jnp.zeros((batch, q_len, kv_len), dtype=jnp.float32)
return mask, indices, score
mask, indices, score = cached_indexer_state
mask = checkpoint_name(mask, "shared_layer_reused_mask")
indices = checkpoint_name(indices, "shared_layer_reused_indices")
return mask, indices, score| and getattr(config, "prune_shared_indexers", True) | ||
| and self.is_shared_layer | ||
| ) | ||
| if self.use_indexer and not is_pruned: |
There was a problem hiding this comment.
If self.use_indexer is True but is_pruned is True (which happens on shared layers when prune_shared_indexers is enabled), the self.indexer attribute is never initialized. However, in __call__ (line 1337), self.indexer is accessed directly as if self.indexer is not None:. This will raise an AttributeError: 'MLA' object has no attribute 'indexer' on shared layers. Initialize self.indexer = None before the conditional block to ensure the attribute always exists.
| if self.use_indexer and not is_pruned: | |
| self.indexer = None | |
| if self.use_indexer and not is_pruned: |
| prune_shared_indexers: bool = Field( | ||
| True, | ||
| description="Whether to prune indexer parameters on Shared (S) layers when use_index_share is enabled.", | ||
| ) |
There was a problem hiding this comment.
When scan_layers=True and prune_shared_indexers=True, the layers in the scanned block (e.g., moe_layers) will have different PyTree structures because some layers will have the indexer parameters while others (the shared layers) will have them pruned. This structural mismatch will cause jax.lax.scan to fail with a PyTree structure mismatch error during tracing. Please add a validation check to prevent prune_shared_indexers=True when scan_layers=True, or ensure that the PyTree structures of the scanned layers are unified/padded.
| inputs: jnp.ndarray, | ||
| position: jnp.ndarray, | ||
| unsqueeze_dim: int | None = 1, | ||
| unsqueeze_dim: int | None = 2, |
There was a problem hiding this comment.
Changing the default value of unsqueeze_dim from 1 to 2 in a shared module like RotaryEmbedding can cause silent shape mismatches or incorrect broadcasting in other models that rely on the default value of 1 (e.g., if they use a different attention layout). It is much safer to keep the default as 1 and explicitly pass unsqueeze_dim=2 when calling RotaryEmbedding in GLM-5.2/MLA.
| unsqueeze_dim: int | None = 2, | |
| unsqueeze_dim: int | None = 1, |
|
|
||
| if getattr(self.config, "use_index_share", False) and cached_indexer_state is not None: | ||
| if layer_idx is not None and self.is_full_tuple is not None: | ||
| is_full = jnp.array(self.is_full_tuple, dtype=jnp.bool_)[layer_idx] |
There was a problem hiding this comment.
Recreating JAX constant arrays like jnp.array(self.is_full_tuple, dtype=jnp.bool_) inside __call__ on every invocation adds unnecessary overhead to tracing and compilation. It is much more efficient to pre-compute these JAX arrays in __init__ and store them as class attributes (e.g., self.is_full = jnp.array(self.is_full_tuple, dtype=jnp.bool_)).
| is_full = jnp.array(self.is_full_tuple, dtype=jnp.bool_)[layer_idx] | |
| is_full = self.is_full[layer_idx] |
| if step is not None: | ||
| return self.start_initial_profile_step <= step <= self.finished_initial_profile_step |
There was a problem hiding this comment.
The is_active(step) method ignores periodic profiling steps because it only checks if step is between start_initial_profile_step and finished_initial_profile_step. This means that during periodic profiling steps, is_active(step) will return False, and train.py will not block/sync metrics, potentially leading to profiling artifacts. Consider checking self._active directly or updating the condition to include periodic profiling steps.
| else: | ||
| tokens = self.data |
There was a problem hiding this comment.
In get_result_at_slot, if self.tokens_idx is None, tokens is set to self.data (the entire batch) instead of being sliced by slot (i.e., self.data[slot]). This is inconsistent with the method's purpose of returning the result for a single slot. Consider slicing self.data by slot when self.tokens_idx is None.
| else: | |
| tokens = self.data | |
| else: | |
| tokens = self.data[slot] if self.data is not None else None |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
f4d987b to
ac4278e
Compare
…ns, and optional elastic imports
154c165 to
ab14b28
Compare
…static pytree mutation
…l carry instead of mutating self
1 participant
Description
This PR introduces native support for GLM-5.2 Cross-Layer Index Sharing (IndexShare) and Indexer Parameter Pruning in MaxText.
In GLM-5.2 (744B), sparse attention is shared across groups of decoder layers following a configurable topology (e.g.
FFFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSacross 78 layers). Only 21 Full (F) layers instantiate indexers and compute query/key projections and Top-K sorting, while the remaining 57 Shared (S) layers reuse the cached indexer state (indexer_mask,topk_indices,indexer_score) with zero projection or sorting overhead.Key Changes and Technical Implementation:
nn.scanandNNXDecoder:NNXDecoder._apply_layers_sequentiallyandscannable_blockto carrycached_indexer_state = (mask, indices, score)across scanned loop iterations.MLA.__call__:jax.lax.condbranching insideMLAto execute_run_fullon Full layers and_run_shared(reusing the mask and indices without compute) on Shared layers.prune_shared_indexers=true):Indexersub-module and its weight parameters on the 57 shared layers, saving substantial memory footprint and eliminating unneeded weight transfers.1 / group_sizeacross shared layer blocks, made fully compatible with dynamic JAX scan tracing.mla_qk_head_chunk_sizeincalculate_indexer_lossusingjax.lax.scanacross head chunks to avoid large intermediate[B, H, T, S]allocations during training.param_mapping.pyto seamlessly convert and load HuggingFace GLM-5.2 checkpoints containing only 21 physical indexer tensors.Tests
seq=2048,per_device_batch_size=1).glm_full_layer_indexerandglm_shared_layer_index_reuse.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.