Repository navigation
0.3.2a3: compile segment-mode cluster locality, recover from a failed compile trace - #55
Merged
Merged
Conversation
… compile trace - cluster_locality_mode="segment" now traces under mx.compile: document ids come from segment_ids_mx(input_ids, chunk=window) and the loss uses static shapes (B*S segment slots, unused slots weigh 0). Matches the NumPy ids and the Torch loss. - CompiledStep snapshots weights, RNG key and side outputs before a first trace and restores them when the trace fails, so the eager fallback no longer trains on tracers. - tests/test_segment_compile.py (5 cases; 4 fail on 0.3.2a2). Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
FonaTech
force-pushed
the
long-context-0.3.2a3
branch
from
November 10, 2026 06:32
4e8dea1 to
e32660d
Compare
MLX does not bounds-check gathers, so an id >= vocab_size read neighbouring memory and crashed the process at random. run_mlx_stage (every batch, GRPO prompts) and LongRun now raise a ValueError naming the id range. test_mlx_training_six_stage_loss_paths crashed about one run in two on CPU, on clean dev as well: its GRPO stage tokenizes with the bundled 6400-token tokenizer while the model had vocab_size=64 (ids up to 6170). The test now uses vocab_size=6400 (25/25 runs pass). Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes three kernel problems found while benchmarking the 50M and 30M multilingual configs on the real corpus.
cluster_locality_mode="segment"could not be compiled: the loss built document ids with NumPy on a traced array.CHRONOS_MLX_COMPILE=onraised, andautocrashed on the nextmx.eval. It now usessegment_ids_mx(input_ids, chunk=window)with static shapes. Ids equal the NumPy rule on 300 random batches; the loss matches the Torch reference within 1.7e-5.CompiledStepnow snapshots and restores them; the RNG stream continues exactly.vocab_sizeare rejected with aValueErrorbefore MLX reads the embedding table. MLX does not bounds-check gathers, so such ids crashed the process at random. This was the cause of the intermittent crash intest_mlx_training_six_stage_loss_paths(about 1 run in 2 on CPU, also on clean dev): its GRPO stage tokenizes with the bundled 6,400-token tokenizer while the model hadvocab_size=64. The test now uses 6400 and passes 25 of 25 runs.Tests:
tests/test_segment_compile.py(5 cases, 4 fail on 0.3.2a2) andtests/test_vocab_range_guard.py(3 cases). Local suite with MLX on CPU: 600 passed, 1 skipped. ruff 0.17.4 clean.🤖 Generated with Claude Code