Fix silent wrong-result lowering; add MobileCLIP2-S2 end-to-end coverage - #37
Merged
Merged
Conversation
Driven by a MobileCLIP2-S2 parity hunt (attention outputs diverged while MLP/norm/downsample matched). Several independent silent-correctness bugs fell out; all are fixed here with coverage in both runners. Correctness fixes (numerics change vs main — all toward the framework reference, so existing exact-GELU / reduction / clamp / SDPA results shift): - GELU: torch exact (approximate='none', the default) lowers to GGML_UNARY_OP_GELU_ERF (16) end to end (importer attribute, lowering, codegen, interpreter); tanh keeps GELU (8). - Single-axis MEAN/SUM over ggml dims 2/3 now swap that axis to position 0 and reduce it in both runners; the old flatten-HxW path reduced the wrong elements. - Multi-axis MEAN/SUM/AMAX/AMIN chain into single-axis ops through the shared ir.graph.chain_single_axis_reduction helper (torch + JAX); lowering rejects unchained multi-axis specs instead of silently using axes[0]. - SDPA without an explicit scale uses 1/sqrt(head_dim) (was 1.0). - CLAMP with an open min/max defaults to -/+FLT_MAX in both runners (was 0/6, silently bounding clamp_min/clamp_max). - Grouped convs with a channel multiplier (groups == in != out) decompose into per-group convs; lowering rejects non-depthwise grouped convs instead of mis-emitting CONV_2D_DW. True depthwise still lowers to one DW op. - JAX TF-SAME convs with asymmetric padding materialize the excess sides as zero concats and keep the symmetric minimum (was: before-side only). - Depthwise convs always use ggml_conv_2d_dw_direct with the mirrored match_dw_layout helper (1D folding, F32 casts) in both runners. Loud failures instead of silent wrong graphs: - Lowering raises NotImplementedError for inexpressible 5D permutes (was: silently a RESHAPE, scrambling QKV head/slot order); the torch importer squeezes static unit dims so QKV-style permutes become plain 4D permutes. - Slices of folded (>4D) dims land on ggml dim 3 with an offset multiplier (was: out-of-range nb[4] in generated code). - Unhandled opcodes in the C++ emitter raise (was: identity passthrough). - Graph arena sized from node count via ggml_new_graph_custom. Compat / plumbing: - GGUF writer shortens tensor names >= GGML_MAX_NAME (64) to name[:54]_sha1[:8], applied consistently to the tensor-info table and the graph-spec JSON the native loader binds weights by. - JAX explicit-pad tensors use padconv_N names (jaxpr brackets/colons broke the GGUF loader); conv bias adds go through match_broadcast (incl. I32->F32 casts); CONCAT supports N inputs and skips empties; flash-attn output is permuted to graph layout unless fused upstream; fused conv+relu; aten.expand_as mapping; unbind/split-container getitem support. - New examples/models/openclip_model.py MobileCLIP2-S2 loader (hf-hub timm checkpoint, reparameterize_model, normalizing wrapper, preprocessor) plus open_clip_torch dev dependency and a CPU e2e parity test. Tests: per-dim MEAN/SUM, chained multi-dim reductions, exact/tanh GELU, QKV 5D permute, SDPA default scale, conv bias/depthwise/grouped-multiplier, clamp_min, expand_as, folded-slice views, GGUF long names — in both the standalone (codegen) and interpreter runners — plus importer/lowering structure tests and JAX reduction/padding tests.
Owner
This is also my way of developing new lowering rules in ggmlc. Choose a target model from a well known architecture family, introduce mappings, fix shape manipulations, add tests so that we can make sure that we don't break anything when adding new rules in the future. This way, each addition will enable to compile some other unseen architectures as well.
Great. I also have some ideas to increase maintainability of these two paths + a better typed codebase, and I hope that will land early this week.
That's correct. If my memory serves me correctly, only the official CLIP model by OpenAI was sensitive to using quick GELU, but I'll double-check it later on. This PR looks good to me. Tests are also passing in all the platforms, merging. Thanks |
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
2 participants
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.
I wanted to compile and run MobileCLIP2-S2, and I found a few unsupported operations such as
unbindandexpand_as, as well as several shape-manipulation issues. I went down a bit of a rabbit hole, and this PR ended up being bigger than originally intended.This time, I've put special care into matching the behavior between codegen and the executor. I've also added tests for many of the changes, including a new end-to-end test for MobileCLIP2-S2 from OpenCLIP.
I think most of the changes are fairly uncontroversial, with the possible exception of changing GELU to use the exact implementation instead of the approximation. I believe this is the correct behavior, since otherwise we cannot reproduce the same results as the PyTorch model.
I've also changed
conv_2d_dwtoconv_2d_dw_direct, which is generally faster for depthwise convolutions.There were a few other issues uncovered along the way:
CLAMPwas incorrectly using[0.0, 6.0]as its default range.The commit message contains the complete list of changes.