Repository navigation
Fix Engram's n-gram hashing and ShortConv under context-parallel load balancing - #5992
Conversation
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
a8507d1 to
cbe5110
Compare
… load balancing When `context_parallel_load_balance=True` (the default), the training input pipeline permutes each sequence into `DUAL_CHUNK_SWAP` order via `max_utils.reorder_sequence`, splitting the sequence into `2 * cp_size` chunks so context-parallel rank `r` holds chunks `r` and `2 * cp_size - 1 - r`. Two operations inside Engram previously indexed into the permuted sequence as if adjacent array positions were adjacent time steps: 1. `NgramHashMapping._get_ngram_hashes` shifted along the array axis (`x[:, :-k]`) to form n-gram windows, pairing tokens at every chunk seam with unrelated history or future tokens. 2. `ShortConv` ran a causal depthwise 1D convolution (`kernel_size=4`, `dilation=3`, receptive field `halo_size = 9`) directly across the permuted chunks, mixing non-adjacent chunks and leaking future tokens from chunk `r` into chunk `2 * cp_size - 1 - r`. This commit fixes both operations without redistributing full activation tensors: - `NgramHashMapping.__call__(input_ids, load_balanced_cp_size=1)`: un-permutes `compressed_ids` (`int32 [B, S]`) back to contiguous natural order before computing `_get_ngram_hashes`, then re-permutes `hash_ids` back to `DUAL_CHUNK_SWAP` order. - `ShortConv._load_balanced_conv`: replaces the `CAUSAL` left zero-padding of `nnx.Conv` with a `jax.shard_map` (manual only over `config.context_sharding`) that exchanges the `halo_size`-token boundary tails from each chunk's natural predecessor via `lax.ppermute`, followed by the same `VALID` depthwise convolution (`lax.conv_general_dilated`). - `DeepSeekGenericLayer.engram_op`: queries `max_utils.reordered_cp_size` in `MODEL_MODE_TRAIN`, raises `ValueError` if a non-`DUAL_CHUNK_SWAP` reorder strategy (such as GPU `STRIPED`) is configured, and passes `load_balanced_cp_size` to `ngram_hash_mapping` and `engram`. - `max_utils.resolve_reorder_strategy`: centralizes `ReorderStrategy.AUTO` resolution (shared with `train_utils.setup_train_loop`).
`CompressedTokenizer` stored its lookup table and `NgramHashMapping` its hash multipliers as `jnp.array`s created in `__init__`. The training loop builds the model under a JAX trace (the jitted `initialize_state`), where those arrays are tracers. The `NgramHashMapping` that holds them is a static part of the model graph that the train step later runs, so every Engram model failed in its first train step with `UnexpectedTracerError` (int32[vocab_size], created in `CompressedTokenizer.__init__`). Keep both as numpy arrays and convert the lookup table when it is used. The values are unchanged.
Add load-balanced variants of the hash and ShortConv tests in engram_vs_reference_test.py. These are the two Engram ops that read neighbouring tokens; everything else is per token. Each test feeds a DUAL_CHUNK_SWAP-permuted sequence, as the training input pipeline does under context_parallel_load_balance, with CP = 2 and 4, and compares against natural order: - test_hash_load_balanced: the hash ids are identical. - test_shortconv_load_balanced: with a non-zero kernel and the sequence sharded over the context axis, outputs and input and parameter gradients match. The file now asks XLA for four host devices, so that the context axis exists on CPU. Like the rest of the file, the tests need torch, and pytest.ini keeps them out of CI; run the file directly. Based on Kuo Wei's tests for the fix.
cbe5110 to
21d6dd0
Compare
|
🤖 Hi @NuojCheng, I'm running an automatic MaxText review of this pull request. You can track my progress in the logs. To request another review later, add the |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses silent correctness issues in the Engram model when training with context-parallel load balancing enabled. By default, the input pipeline reorders sequences, which previously caused Engram's n-gram hashing and causal convolution to read incorrect tokens at chunk boundaries. The changes introduce logic to un-permute inputs for these operations and implement a cross-rank token exchange for the convolution, ensuring identical outputs to natural-order processing. Additionally, the PR fixes a JAX tracer error that prevented model initialization on main. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
📋 Review Summary
Verdict: Spec met; correctly fixes Engram under CP load balancing and resolves the tracer leak, but breaks when engram_kernel_size=1.
Top issues:
- [Bug] If
engram_kernel_size=1,halo_sizebecomes0. The Python slicefirst[:, -halo:]then evaluates tofirst[:, 0:](the entire chunk), which doubles the sequence length inside_load_balanced_convand breaks shapes. → Add an early return forhalo == 0since a pointwise convolution commutes with the sequence permutation (see inline comment inengram.py).
Inferred spec & validation
- Sources: PR description; code before/after.
- Problem: Engram reads the wrong history tokens under
context_parallel_load_balance=Truedue to sequence permutation, andtrain.pycrashes due to a JAX tracer leak inCompressedTokenizer/NgramHashMapping.
| Spec item | Implementation | Validating test |
|---|---|---|
| R1 (stated): Un-permute tokens before hashing and permute hash ids back. | ✅ | test_hash_load_balanced ✅ |
R2 (stated): Replace CAUSAL padding in ShortConv with natural predecessors via ppermute. |
✅ | test_shortconv_load_balanced ✅ |
R3 (stated): Resolve context_parallel_reorder_strategy=auto consistently. |
✅ | ❌ none |
R4 (stated): Convert jnp.array in __init__ to NumPy arrays to avoid tracer leak. |
✅ | ❌ none |
| C1 (inferred): Existing checkpoints and configs are unchanged. | ✅ preserved | ❌ none |
| conv = self.conv | ||
| if seq_len % (2 * cp_size) != 0 or seq_len // (2 * cp_size) < halo: |
There was a problem hiding this comment.
🟠 If engram_kernel_size is configured to 1, halo becomes 0. The Python slice first[:, -0:] evaluates to first[:, 0:] (fetching the entire chunk instead of an empty slice). This causes halo_first to have length L instead of 0, doubling the sequence length inside the convolution and breaking shapes. An early return avoids this and correctly applies the pointwise projection directly, since a pointwise operation commutes with the sequence permutation.
| conv = self.conv | |
| if seq_len % (2 * cp_size) != 0 or seq_len // (2 * cp_size) < halo: | |
| conv = self.conv | |
| if halo == 0: | |
| return conv(x_flat) | |
| if seq_len % (2 * cp_size) != 0 or seq_len // (2 * cp_size) < halo: |
There was a problem hiding this comment.
nit: +1 — confirmed on CPU: ShortConv(..., kernel_size=1, dilation=3, mhc_expansion_rate=4)(x, load_balanced_cp_size=2) on x of shape (2, 16, 4, 16) yields halo == 0, turns first[:, -halo:] into first[:, 0:], and raises TypeError: cannot reshape array of shape (2, 32, 64) (size 4096) into shape (2, 16, 4, 16) (size 2048) at ShortConv.__call__.
Simplest fix is gating directly at ShortConv.__call__ (if load_balanced_cp_size > 1 and self.halo_size > 0:), which skips _load_balanced_conv (including its 2 * cp_size divisibility check, shard_map, and both ppermutes) entirely when the convolution is width-1.
There was a problem hiding this comment.
Fixed with the gate you suggested: ShortConv.__call__ takes the halo path only when load_balanced_cp_size > 1 and self.halo_size > 0. engram_test.py covers a width-1 kernel under CP=2.
shuningjin
left a comment
shuningjin
left a comment
There was a problem hiding this comment.
Thanks for the fix! LGTM on the CP load-balancing logic, partial shard_map halo exchange, and JIT tracer-leak fix.
| np.testing.assert_equal(pt_out, jax_out) | ||
|
|
||
| @parameterized.named_parameters(("cp2", 2), ("cp4", 4)) | ||
| def test_hash_load_balanced(self, cp_size): |
There was a problem hiding this comment.
Neither test_hash_load_balanced nor test_shortconv_load_balanced runs in CI because pytest.ini passes --ignore=tests/unit/engram_vs_reference_test.py (this file unconditionally imports torch and downloads deepseek-ai/DeepSeek-V3 at test_hash_load_balanced, leaving patch coverage at 41.7%). In addition:
- There is no regression test instantiating
NgramHashMapping/CompressedTokenizer.lookup_tableinsidejax.jitto guard against theUnexpectedTracerErrorfixed in commit 2. os.environ.setdefault("XLA_FLAGS", ...)is skipped wheneverXLA_FLAGSis already set in the shell without--xla_force_host_platform_device_count, causingget_cfg_and_mesh(..., ici_context_parallelism=cp_size)to raise on a 1-device host instead of skipping cleanly.
Suggestion: Move the pure-JAX CP load-balancing checks (plus a jax.jit construction test for NgramHashMapping and a max_utils.resolve_reorder_strategy check) into a standalone tests/unit/engram_cp_test.py with no torch import and the checked-in src/maxtext/assets/tokenizers/qwen3-tokenizer (or a stub HF tokenizer); cpu-unit CI already runs with XLA_FLAGS=--xla_force_host_platform_device_count=4 (run_tests_against_package.yml), and guarding with if jax.device_count() < cp_size: self.skipTest(...) handles direct 1-device invocations.
Follow-up OK: Can be addressed in a follow-up PR so #5992 merges unblocked; without a non-ignored pure-JAX unit test, future regressions in _load_balanced_conv, load-balanced hashing, or the JIT tracer-leak fix will not fail CI.
There was a problem hiding this comment.
Thanks, addressed in this PR rather than in a follow-up:
- The load-balancing tests moved to a torch-free
tests/unit/engram_test.py. It uses the checked-in qwen3 tokenizer, so it runs incpu-unit, where all 6 tests pass; patch coverage is now 84%. The ShortConv cases skip when there are fewer devices thancp_size. test_built_under_jitbuildsNgramHashMappingunderjax.jitand runs it in a later jit. It fails withUnexpectedTracerErrorif the tables go back to jax arrays.max_utils_test.pyhas a newresolve_reorder_strategytest.engram_vs_reference_test.pyis back to its state on main.
One note on the ShortConv test: it shards the batch on fsdp, as training does. With a replicated batch and a sharded sequence, XLA scales nnx.Conv's kernel gradient by the replica count. That is the existing non-load-balanced path, which this PR does not change.
| conv = self.conv | ||
| if seq_len % (2 * cp_size) != 0 or seq_len // (2 * cp_size) < halo: |
There was a problem hiding this comment.
nit: +1 — confirmed on CPU: ShortConv(..., kernel_size=1, dilation=3, mhc_expansion_rate=4)(x, load_balanced_cp_size=2) on x of shape (2, 16, 4, 16) yields halo == 0, turns first[:, -halo:] into first[:, 0:], and raises TypeError: cannot reshape array of shape (2, 32, 64) (size 4096) into shape (2, 16, 4, 16) (size 2048) at ShortConv.__call__.
Simplest fix is gating directly at ShortConv.__call__ (if load_balanced_cp_size > 1 and self.halo_size > 0:), which skips _load_balanced_conv (including its 2 * cp_size divisibility check, shard_map, and both ppermutes) entirely when the convolution is width-1.
| # history and the causal convolution are the two places Engram looks at neighbouring tokens; | ||
| # both take the permutation into account. Inference (prefill / autoregressive) never reorders. | ||
| load_balanced_cp_size = 1 | ||
| if self.model_mode == MODEL_MODE_TRAIN: |
There was a problem hiding this comment.
nit: DeepSeekGenericLayer.engram_op checks construction-time self.model_mode == MODEL_MODE_TRAIN, whereas both callers (DeepSeekDenseLayer.__call__ and DeepSeekMoELayer.__call__) take a call-time model_mode so a single NNX model instance can be called across modes (the invariant documented in DecoderLayerModelModeTest). If a MODEL_MODE_TRAIN-constructed model with context_parallel_load_balance=True is called with MODEL_MODE_PREFILL or MODEL_MODE_AUTOREGRESSIVE, engram_op will still un-permute/halo-exchange non-reordered tokens (or raise ValueError on S = 1). Pass call-time model_mode into engram_op(self, x, decoder_input_tokens, model_mode) and gate on model_mode == MODEL_MODE_TRAIN instead.
There was a problem hiding this comment.
Good catch, fixed. engram_op now takes the call-time model_mode from DeepSeekDenseLayer.__call__ and DeepSeekMoELayer.__call__, and gates load balancing on it.
With engram_kernel_size=1 the halo is empty, and `first[:, -0:]` slices the whole chunk instead of nothing. The load-balanced convolution then doubled the sequence length and failed to reshape. A width-1 kernel reads no history, so call self.conv directly whenever halo_size is 0.
engram_op checked the model_mode the layer was built with, but one NNX model serves train, prefill and autoregressive calls. A model built for training and called for inference still un-permuted the tokens and exchanged conv halos, even though inference never reorders the sequence. Pass the call-time model_mode in from both DeepSeek layers.
engram_vs_reference_test.py needs torch and a tokenizer download, so pytest.ini keeps it out of CI, and the load-balancing tests never ran there. Move them to tests/unit/engram_test.py. That file uses the checked-in qwen3 tokenizer and runs on the 4 CPU devices of the cpu-unit job. The ShortConv cases skip when there are fewer devices than context-parallel ranks. The input batch is now sharded on fsdp, as in training, because XLA scales nnx.Conv's kernel gradient by the replica count when the batch is replicated while the sequence is sharded. Also add: - a width-1 kernel case to the ShortConv test; - a test that builds NgramHashMapping under jit and runs it in a later jit, which fails with UnexpectedTracerError without the numpy tables; - a resolve_reorder_strategy test in max_utils_test.py. engram_vs_reference_test.py is back to its state on main.
Description
With
context_parallel_load_balance=True(the default), the training input pipeline puts every sequence in DUAL_CHUNK_SWAP order (max_utils.reorder_sequence): the sequence is cut into2 * cpchunks, and context-parallel rankrgets chunksrand2 * cp - 1 - r. Engram reads each token's history by array index in two places. With CP > 1 it therefore used the wrong tokens at every chunk seam, some of them from later in the sequence:NgramHashMapping._get_ngram_hashesbuilds its windows withx[:, :-k]. The firstmax_ngram_size - 1tokens of every chunk except chunks0andcphashed the wrong n-grams, so they read the wrong memory rows.max_ngram_size= 3, so 9 tokens of history) ran straight across the permuted chunks.The rest of Engram is per-token. Nothing rejected this combination, so any Engram model trained with CP > 1 got silently wrong outputs, and its loss still went down.
Sections 1–3 below describe the first three commits. The last three commits address review; they are described in section 4.
1. The fix (Kuo Wei)
NgramHashMapping.__call__(input_ids, load_balanced_cp_size)un-permutes the compressed token ids (int32[B, S]), hashes them in natural order, and permutes the hash ids back.ShortConv._load_balanced_convreplaces theCAUSALzero padding with the last 9 tokens of each chunk's natural predecessor, then runs the sameVALIDconv thatnnx.Convruns. The tokens come from twolax.ppermutes inside ashard_mapthat is manual only overcontext_sharding. Each rank sends 2 × 9 tokens per sequence, whatever the sequence length; no activation is all-gathered or un-permuted.DeepSeekGenericLayer.engram_openables both in train mode whenmax_utils.reordered_cp_size(config, mesh) > 1, the helper the LoadBalanced splash masks and MTP already key off. It raises for a reorder other than DUAL_CHUNK_SWAP (GPU STRIPED), andShortConvraises if a chunk is shorter than the conv's receptive field.max_utils.resolve_reorder_strategyresolvescontext_parallel_reorder_strategy=autoin one place, for bothsetup_train_loopandengram_op.2. Engram crashed in
train.pyon main, with or without CPOn main,
train.pyfails before the first step for every Engram model:CompressedTokenizer(lookup table) andNgramHashMapping(hash multipliers) storedjnp.arrays created in__init__, and the training loop builds the model under a JAX trace. Both are now numpy arrays; the lookup table is converted where it is used, and the values are unchanged.train_compilewas not affected, which is why AOT tests did not catch it. This commit is needed to train Engram at all, and so to test commit 1 end to end.3. Tests
tests/unit/engram_test.pyneeds neither torch nor a download: it uses the checked-in qwen3 tokenizer. CI'scpu-unitjob runs it on 4 CPU devices. The tests cover the two ops that read neighbouring tokens; the rest of Engram is per-token. Each test feeds a DUAL_CHUNK_SWAP-permuted sequence and compares the result against natural order:NgramHashMappingTest.test_load_balanced(CP 2 and 4): the hash ids are identical.ShortConvTest.test_load_balanced(CP 2 and 4, plus a width-1 kernel): with a non-zero kernel and the sequence sharded over the context axis, outputs and input and parameter gradients match. These cases skip when there are fewer devices than CP ranks.fsdp, as in training. With a replicated batch and a sharded sequence, XLA scalesnnx.Conv's kernel gradient by the replica count. That affects the existing non-load-balanced path and is not changed here.NgramHashMappingTest.test_built_under_jitbuilds the hash mapping underjitand runs it in a laterjit. Without commit 2 it fails withUnexpectedTracerError.max_utils_test.pygets aresolve_reorder_strategytest.4. Review follow-ups
engram_kernel_size=1the halo is empty, andfirst[:, -0:]took the whole chunk, so the load-balanced conv failed to reshape.ShortConvnow uses the plain conv whenhalo_sizeis 0; a width-1 kernel reads no history.engram_opchecked the mode the layer was built with. One NNX model serves train and inference calls, so it now gates on themodel_modepassed to the layer's__call__.engram_vs_reference_test.py, which pytest.ini keeps out of CI, to the torch-freeengram_test.py(section 3).engram_vs_reference_test.pyis unchanged from main.Tests
Unit, CPU.
engram_test.py: all 6 tests pass with 4 devices (thecpu-unitsetting) and with 8. With 1 device the 3 ShortConv cases skip. Each test fails when the fix it covers is reverted:test_built_under_jitfails;halo_sizegate removed: the width-1 case fails.The tolerance is
rtol=atol=1e-4. Outputs match exactly, input gradients agree within 9.5e-7, and parameter gradients within 1.6e-5 on values up to 72.End to end,
train.pyon CPU with 8 devices. A smalldeepseek-custom: 3 layers, Engram on layer 1, mHC 4, MLA withattention=dot_product, indexer off, float32 withmatmul_precision=highest, synthetic data, global batch 8 × 128. Losses are computed from the logged perplexity:With this PR, load balance on matches load balance off to every printed digit. Without the fix it is off by 1.7e-4 already at step 0. The conv kernel is zero-initialized, so step 0 shows only the hashing error. The ~1e-6 between CP=1 and CP=2 comes from sharding. On main without commit 2, all of these runs crash with the tracer error above.
Command
AOT, v5p-8. Full
deepseek-custom(Engram on layers 1 and 4),ici_context_parallelism=2, load balance on,per_device_batch_size=2,max_target_length=1024. It compiles both withattention=flash use_tokamax_splash=False use_indexer=Falseand withattention=dot_product, at 15.89 / 15.90 GiB peak. Tokamax splash was not tried because this environment's tokamax predates main.Existing tests.
max_utils_test.py,train_utils_nnx_test.py,mhc_test.pyandengram_test.pypass (79 passed, 3 skipped).decoder_layer_model_mode_test.pypasses (32).engram_vs_reference_test.pyis not changed by this PR. On main,test_hash_matchandtest_engram_match_*fail: since 3138907, Engram draws its hash multipliers with int32 bounds while the PyTorch reference uses int64.Not run on TPU hardware.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.