feat(xla): support Molmo sparse additive visual embeddings - #918
feat(xla): support Molmo sparse additive visual embeddings#918inureyes wants to merge 16 commits into
Conversation
inureyes
added
area:inference
ac54a07 to
2c6c9d4
Compare
|
Bounded actual-checkpoint follow-up on commit 6481443: all pinned processor boundaries passed exactly (677,376 patch values, 1,152 mask values, image token IDs, sparse image_input_idx, and grid metadata). The staged eager capture then materialized the patch embedding successfully and entered selected ViT layer 14. That named stage produced no completion marker within the five-minute no-output limit and was interrupted once with exit 130. The CPU run will not be repeated unchanged; the remaining numeric oracle requires a small GPU eager diagnostic runner or a cached/staged equivalent. No parity result is claimed, and the PR remains draft. |
2dc2866 to
d20b06e
Compare
|
Rebased onto Post-rebase validation:
The Molmo runtime now uses the explicit |
Preserve processor-supplied image_input_idx row correspondence across negative sentinels and reject duplicate, out-of-range, over-capacity, malformed, or non-finite inputs before native execution. Keep patch-mask rows distinct from pooled feature rows and prove an embedding-replacement negative fixture differs from the required additive result. Refs #870
Preflight every destination sum so non-finite overflow cannot partially mutate a reusable batch slot, and reject processor mask values outside the qualified [-1,1] coverage-and-sentinel domain. Keep raw patch-mask rows distinct from pooled feature rows and split tests out of the production module to preserve file-size limits. Refs #870
Reuse the shared XLA dense language graph for Molmo v1 by binding its OLMo-style checkpoint names, interleaved RoPE, fused QKV bias, and fused SwiGLU projections without constructing a second decoder. Interpret Molmo intermediate_size as the combined gate/up width and test every fused row slice against the checkpoint schema. Refs #870
Load the pinned Molmo v1 vision, pooling, and projector weights into one auxiliary IREE module, retain only the MLX text embedding table for prepared prefill, and route CLI/server preprocessing through an observable backend policy. Bind static dimensions and OLMo architecture into the artifact identity, validate allocation arithmetic before invocation, and cover the replacement-vs-addition logit sentinel. Refs #870
Represent Molmo's rank-one fused QKV bias with a distinct element-slice schema instead of passing it through the rank-two projection row slicer. Reject incompatible Molmo architecture variants before graph construction so the pinned OLMo layout cannot silently drift. Refs #870
Map Molmo's fused FFN value and gate halves into the shared emitter's gate and value argument order so SiLU is applied to the correct half. Perform sparse visual addition in the original F16 or BF16 text-embedding dtype before widening the owned prepared payload to F32. Validation: cargo test -p mlxcel-xla molmo --lib; cargo test -p mlxcel-xla phi4 --lib; cargo clippy --fix --features xla-iree --lib --allow-dirty. Refs #870
The existing real-checkpoint gate compared only the final PreparedPrefill values, so processor, vision, projector, sparse-merge, and decoder propagation regressions could not be localized. Add a CPU local-task reference feature and a diagnostics-only Molmo graph that captures pinned processor contracts, patch embeddings, selected vision layers, projected features, prepared sparse-add rows, and one resident decoder's KV and logits boundary. Negative controls reject Molmo2 token scanning and replacement semantics without loading a duplicate decoder. Validation: Molmo unit tests 15/15; CPU diagnostic tests 2/2 with the actual-checkpoint test ignored; diagnostic StableHLO compiled for IREE local-task; feature clippy completed with pre-existing unrelated warnings. The pinned 5.3 GB actual-checkpoint gate was not run in this task environment. Refs #870
Materialize the eager MLX patch embedding, each selected ViT layer, and the pool/projector at explicit diagnostic boundaries so the pinned #870 gate reports progress before and after every potentially long stage. Retain only configured selected hidden states instead of every ViT layer, bounding the lazy graph and allowing the next actual checkpoint run to identify the first stalled or divergent stage under the five-minute no-output policy. Validation: cargo check --features xla-reference-diagnostics --lib; cargo test --features xla-reference-diagnostics --lib strict_oracle_rejects_molmo2_token_scan_and_replacement; cargo fmt --all -- --check; git diff --check. The actual full-checkpoint gate was intentionally not rerun. Refs #870
Expose the pinned Molmo reference gate as a dedicated example so the eager oracle can run on MLX CUDA while IREE remains on local-task or local-sync. Emit flushed stage progress and periodic heartbeats without changing the existing comparison thresholds.
Keep the rebased Molmo runtime on the explicit legacy artifact path until the real-checkpoint boundary oracle completes.\n\nRefs #870
Rebasing the Molmo prepared-prefill contract onto the numeric-oracle foundation exposed that its F16 and BF16 merge rounding called through a weight-loader module compiled only for IREE or tests, which broke the default mlxcel-xla library build. Move the three scalar conversion helpers into a feature-neutral module while retaining the weight-loader re-exports and existing bit-contract tests. Preserve both diagnostic feature families during the rebase, and document the pinned revision override required by the locally complete checkpoint fixture. Validation: default mlxcel-xla Clippy, 15 focused Molmo tests, four F16 conversion tests, formatting, and diff checks. Refs #870
d20b06e to
a4047b2
Compare
|
Rebased #918 onto The actual-run assets are present and pinned: the 5.0 GB checkpoint, image and config hashes, GB10 CUDA device, and local-task IREE distribution all match. The local fixture lacks Hugging Face metadata, so the runner documentation now supplies its already-supported pinned revision environment value. One Validation completed: default |
|
Additional bounded linker audit after the rebase evidence:
This does not alter the qualification result: #918 remains draft/fail-closed pending a reproducible cached release runner and the original intermediate/KV/token/lifecycle gates. |
The Molmo v1 reference gate created its production IREE preprocessor without first applying the diagnostics-only local-task thread configuration added in #945. `configure_diagnostic_local_task_threads` had no callers on this branch, while the Gemma3 runner already applies it, so the Molmo runner aborted before any comparison ran: xla_aux_create failed (status 13): iree/base/threading/thread_pthreads.c:159: INTERNAL; thread creation failed with 22 IREE parses its process-global flag registry when the first instance is created, so the call has to precede `load_iree` to have any effect. With this, the pinned 5.0 GB checkpoint gate completes end to end for the first time on GB10: processor stages exact, patch embedding, selected layer 14, projector, both prepared sparse-add merge stages, layer-0 KV, all-layer KV and prefill logits all pass, leaving 4 of 1,181,696 values outside tolerance at selected layer 21. Refs #870
First completed actual for this gate, at
|
| stage | status | elements | max_abs | failures |
|---|---|---|---|---|
processor.patches |
PASS | 677376 | 0.0 exact | 0 |
processor.patch_masks |
PASS | 1152 | 0.0 exact | 0 |
processor.image_token_ids |
PASS | 316 | 0 | |
processor.image_input_idx |
PASS | 288 | 0 | |
processor.grid |
PASS | 3 | 0 | |
vision.patch_embedding |
PASS | 1179648 | 1.336575e-3 | 0 |
vision.selected_layer_14 |
PASS | 1181696 | 2.258102 | 0 |
vision.selected_layer_21 |
FAIL | 1181696 | 2.317841 | 4 |
vision.projector |
PASS | 1032192 | 2.251205e-1 | 0 |
prepared.non_visual_rows |
PASS | 132608 | 0.0 exact (atol=0, rtol=0) | 0 |
prepared.sparse_add_visual_rows |
PASS | 1032192 | 2.500000e-1 | 0 |
decoder.layer0_kv |
PASS | 128 | 0.0 exact | 0 |
decoder.all_layer_kv |
PASS | 3584 | 5.212402e-2 | 0 |
decoder.prefill_logits |
PASS | 152064 | 1.176004e-1 | 0 |
Exit 101 is the intended first-divergence panic on the single failing stage.
Four values out of 1,181,696, and the divergence does not propagate: the projector that
consumes the ViT output, the authoritative sparse-add merge, all-layer KV and the prefill
logits all pass downstream of it. elements=1181696 is 2 crops times 577 tokens times
1024 hidden, so flat index 514613 is crop 0, token 502, channel 565.
The 4 values are deterministic, and the gate is cross-device by construction
Re-running with MLXCEL_MOLMO_IREE_DEVICE=local-sync, which is single-threaded and
deterministic, reproduces the failure identically: same max_abs=2.317841, same
max_rel=5.661761e3, same 4 failures, same first_failure=514613. So this is not a
local-task worker or reduction-order artifact.
MLXCEL_MOLMO_IREE_DEVICE=cuda is rejected by the gate itself
(the #870 reference gate requires IREE local-task or local-sync), so the comparison is
by design MLX on CUDA against IREE on CPU. The two selected layers straddle the threshold
narrowly under atol + rtol*|expected|:
selected_layer_14max difference 2.258102 passes, its worst element sitting where
|expected|is at least about 40.2selected_layer_21max difference 2.317841 fails, its worst element sitting where
|expected|is below about 41.4
Both layers differ by roughly the same amount; layer 21's worst element simply lands at a
slightly smaller |expected|. That looks like the residual of a GPU reference against a
CPU candidate rather than a semantic mismatch, especially since every downstream production
stage including the prefill logits passes. Deciding it properly needs the comparison to run
with both sides on the same device, which this gate currently forbids.
No tolerance was changed and no production arithmetic was touched. Greedy token exactness,
CLI and server paths, cancellation and slot reuse, and mixed batching were not reached
because the strict first-divergence panic stands. This PR stays draft.
Correction to the numbers above: the reference was TF32, not F32The results I posted earlier were collected with MLX's default
True F32 removes about 12 percent of the drift at both selected ViT layers and 25 to 27 It does not close the gate. The same four elements at flat index 514613 still exceed Two further observations on that residual, both from runs at this head:
Filed the reference-policy issue and the two |
1 participant
Summary
ff_proj[value | gate]halves into the shared emitter[gate, value]argument order so SiLU is applied to the correct halfValidation
cargo fmt --all -- --checkgit diff --checkcargo check --features xla-iree --libcargo clippy --fix --features xla-iree --lib --allow-dirtycargo test -p mlxcel-xla molmo --lib: 14 passedcargo test -p mlxcel-xla phi4 --lib: 9 passed, 5 ignoredcargo check --features xla-reference-diagnostics --libafter the eager stage splitcargo test --features xla-reference-diagnostics --lib strict_oracle_rejects_molmo2_token_scan_and_replacement: 1 passedCUDA_ERROR_NO_DEVICE; a separate escalated root CUDA run reached the production CLIGenerating...boundary, then produced no token within the five-minute no-output limit and was interrupted once without retryDraft limitations
Generating...but returned no token within five minutes; it exited 130 after one bounded interruption and was not retried6481443bnow retains only selected hidden states and materializes patch embedding, selected layer 14, selected layer 21, and pool/projector at separate flushed progress boundaries, but actual cross-oracle parity remains to be collected in the next single bounded runPrepared-prefill parity gate
Commit
d90b64e8adds an ignored real-checkpoint host-MLX vs IREE PreparedPrefill comparison. It derives authoritative visual rows from processor indices, requires exact non-visual rows and request metadata, and checks visual rows with a documented max-abs/cosine/normalized-RMSE envelope. The fixture builds and its prompt, gate/value, and sparse-add dtype regressions pass. A single local-task checkpoint run was interrupted after the five-minute no-output bound before either parity metrics or a result were emitted; it was not retried. This PR therefore remains draft.Staged eager diagnostic
Commit
6481443bremoves the diagnostic retention of every intermediate ViT hidden state. The next actual checkpoint run reports begin/completion markers around patch embedding, the two configured selected layers in execution order, and pool/projector evaluation, so the five-minute no-output policy can stop at a named stage and a completed run can proceed directly to first-divergence comparison. The full checkpoint gate was intentionally not rerun in this commit.Dedicated CUDA eager boundary runner
Commit
c6c8f957exposes the existing pinned Molmo boundary comparison as the standalonexla_molmo_reference_checkexample. The eager reference is forced onto MLX CUDA while IREE remains explicitly pinned tolocal-taskorlocal-sync. Flushed stage markers and configurable 30-second heartbeats cover patch embedding, each selected ViT layer, and pool/projector materialization; the comparison stages andEXACT/VISIONtolerances are unchanged.Build from the repository root:
Run the pinned gate:
The standalone Rust/IREE feature boundary passed
cargo check --example xla_molmo_reference_check --features xla-reference-diagnosticsandgit diff --check. The CUDA release build was intentionally stopped during the shared MLX CUDA C++/NVCC rebuild after resource contention with other epic worktrees was identified; the actual 5.3 GB checkpoint gate was not rerun. The PR remains draft pending one bounded run of the command above.2026-07-27 rebase and runner audit
origin/mainatec6fbb6d; the only conflicts preserved the mergednumeric_ops/xla-micro-oraclefoundation alongside the Molmo modules and device-neutral reference featurea4047b22restores default-library compilation by moving Molmo's scalar F16/BF16 merge rounding out of the IREE/test-only weight-loader module while retaining the existing bit-contract testsMLXCEL_MOLMO_REVISION=5c04b3a418979597b1968e41414ad799c87533e8sm_121arelease build ofxla_molmo_reference_checkcompiled its dependencies but did not finish the final GNU root build/link within the five-minute bound; it was interrupted once with exit 130, produced no executable, and this host has nolldormoldalternativemlxcel-xlaClippy passed; focused Molmo tests passed 15/15; F16 conversion tests passed 4/4; formatting and diff checks passedNo actual checkpoint comparison ran at this head, so intermediate, KV/logit, greedy-token, and lifecycle acceptance remain open. The PR stays draft and the unqualified artifact marker remains unchanged.
Closes #870