Gemma4e2b jax.checkpoint & remat fix - #4847
Conversation
![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 refactors the layer application in _apply_gemma4_small_layers to use jax.lax.scan and jax.checkpoint for rematerialization. Feedback on these changes includes removing an unused variable extract_donor_kv, moving the inline import of xla_metadata outside of the loop, and hoisting loop-invariant computations (global_remat_policy and prevent_cse) outside of the decoder layers loop to avoid redundant overhead.
| donor_k, donor_v = layer.compute_shared_kv(y, decoder_positions) | ||
| shared_kv_states[lyr] = (donor_k, donor_v) | ||
| shared_key, shared_value = donor_k, donor_v | ||
| extract_donor_kv = is_donor |
There was a problem hiding this comment.
The variable extract_donor_kv is assigned but never used anywhere in the function. It can be safely removed.
| shared_key=shared_key, | ||
| shared_value=shared_value, | ||
|
|
||
| from jax.experimental import xla_metadata |
There was a problem hiding this comment.
The import from jax.experimental import xla_metadata is executed in every iteration of the loop. Moving it to the top of the file (or before the loop) is cleaner and avoids redundant import lookups in each iteration.
| global_remat_policy = self.get_remat_policy() | ||
| offload_names = maxtext_utils.get_save_and_offload_names(cfg) | ||
| if offload_names[0] or offload_names[1]: | ||
| save_names, offload_to_device = offload_names | ||
| global_remat_policy = jax.checkpoint_policies.save_only_these_names(*(save_names + offload_to_device)) | ||
|
|
||
| prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg) |
There was a problem hiding this comment.
The computations for global_remat_policy and prevent_cse are loop-invariant as they only depend on self and cfg, which do not change across iterations. Computing them inside the loop for lyr in range(cfg.num_decoder_layers): introduces redundant overhead for all 35 layers.
Please move these definitions outside and before the loop (e.g., before line 2245):
global_remat_policy = self.get_remat_policy()
offload_names = maxtext_utils.get_save_and_offload_names(cfg)
if offload_names[0] or offload_names[1]:
save_names, offload_to_device = offload_names
global_remat_policy = jax.checkpoint_policies.save_only_these_names(*(save_names + offload_to_device))
prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg)1 participant
Description
When training the Gemma 4 Small (E2B/E4B) model with NEW_MODEL_DESIGN=1, XLA LICM was co-scheduling the unrolled 35 decoder layers concurrently, causing the peak HBM overhead during the backward pass to spike past 400GB+ for per device for max_target_length=24576 & per_device_batch_size=8
baseline config & error log
Fix:
Modified _apply_gemma4_small_layers in src/maxtext/layers/nnx_decoders.py to securely wrap both the layer's forward pass and its KV extraction (compute_shared_kv) within a jax.checkpoint boundary.
Plumbed this boundary through a native jax.lax.scan(length=1) block armed with skip-simplify-while-loops_trip-count-one, isolating the compiler graph identically to how the standard _scan_global_layer operates in gemma4.py.
This optimization safely forces peak memory back down to ~75-90GB.
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
Tests
Post-fix config & log
logits check
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.