fix(lora): make adapter weight fusion Conv1D-layout aware

[inureyes]

Summary

LoRA weight fusion added a delta to a base tensor with no shape check and no notion of how the base was stored on disk. Deltas are always [out_features, in_features]; raw HuggingFace GPT-2 exports store their projections in the Conv1D layout [in_features, out_features], and Gpt2Layout only transposes them at model construction, which runs after fusion on the load_model_with_adapter path. Fusion is now layout-aware and shape-checked, and it reports rather than guesses.

The two failure modes this closes

The square attention c_proj ([768, 768] against a [768, 768] delta) broadcast cleanly, so the model silently ran W + Dᵀ instead of W + D. That is a different rank-r update of the same shape: no error, no warning, no shape anomaly, just wrong numbers. This is the case the issue was filed for, and it is the reason a shape check alone is not a fix.

The three non-square projections (c_attn, c_fc, MLP c_proj) could not broadcast, and mlxcel_core::add is not declared fallible in the cxx bridge, so the MLX C++ exception was an uncatchable std::terminate that killed the process instead of returning an error.

What changed

• src/lora/loader.rs: reject_conv1d_layout_fusion takes a whole-map layout verdict before the first add. The verdict comes from the fused QKV projection, which is [n_embd, 3 * n_embd] on disk and [3 * n_embd, n_embd] once transposed and can never be confused; that sibling is the only evidence available, since the square c_proj carries no layout signal of its own. Quantized projections are skipped on the same .scales probe Gpt2Layout::detect uses, and the reported evidence key is the lexicographically smallest match so the diagnostic does not depend on HashMap order.
• src/lora/loader.rs: a shape guard compares base and delta before every add and returns an anyhow::Error naming the adapter layer, the base weight key and both shapes, with an added note when the two are exact transposes. This covers any future family that lands with a transposed on-disk layout.
• Nothing is auto-repaired. Transposing the delta would fix the non-square projections and leave the square one exactly as wrong, because a square weight carries no signal about which orientation was intended. Rejection is the only honest outcome, so the error says what to use instead.
• Both guards live in the shared primitive, so they cover the single-process path through apply_lora_adapters and the pipeline-parallel apply_stage_lora_adapter, which never goes through load_model_with_adapter and would have missed a fix applied at the loading layer.
• Both diagnostics are returned errors, not tracing::warn!, which is a no-op in the mlxcel CLI binary because only src/server/startup.rs installs a subscriber.
• Rejection is per-target, not blanket: an adapter that lands on nothing transposed (an embedding-only adapter, say) still fuses on a Conv1D checkpoint.
• src/lora/loader_tests.rs (new): the inline mod tests moved out to a sibling _tests.rs per the project convention, keeping loader.rs under the file-size limit, plus the new coverage.
• src/models/gpt2.rs, src/models/gpt2_tests.rs: the checkpoint-layout fixtures are now pub(crate) so the LoRA tests assert against the same maps instead of re-deriving them.

Test plan

The layout tests fuse an mlxcel-convention (.lora_a / .lora_b) adapter against the raw HuggingFace GPT-2 fixtures from src/models/gpt2_tests.rs, reused rather than reinvented: a locally invented fixture would assume the very thing under test, whereas raw_hf_weights is the map Gpt2Layout::detect is separately proven to classify as a Conv1D export. Every layout test re-asserts that classification before touching fusion.

• conv1d_checkpoint_rejects_fusion_into_the_square_attention_c_proj asserts base and delta shapes are equal before asserting the rejection, which is what demonstrates that a shape check alone cannot catch this case, and asserts the base tensor is left byte-unchanged.
• conv1d_checkpoint_rejects_every_non_square_projection_without_aborting covers c_attn, c_fc and MLP c_proj. Reaching the assertions at all is the evidence: on an unguarded build each of these aborts the whole test binary with SIGABRT.
• an_already_transposed_gpt2_checkpoint_still_fuses_correctly and conv1d_checkpoint_still_fuses_an_adapter_that_targets_nothing_transposed are positive controls asserting exact fused sums.
• stage_local_fusion_rejects_a_conv1d_layout_stage_weight_map and stage_local_fusion_still_applies_to_an_out_in_stage_weight_map cover apply_stage_lora_adapter end to end through a real on-disk adapter directory.
• a_transposed_base_weight_is_reported_instead_of_reaching_mlx_add and a_mismatch_that_is_not_a_transpose_omits_the_transpose_hint cover the generic guard on a family with no Conv1D signal.
• DEVELOPER_DIR=... cargo test --release --lib --features metal,accelerate lora (20 passed)
• DEVELOPER_DIR=... cargo test --release --lib --features metal,accelerate partial_loading_adapter (14 passed, including the pre-existing stage_filtered_fusion_matches_full_fusion_on_stage_subset parity test, so [out, in] behavior is unchanged)
• DEVELOPER_DIR=... cargo test --release --lib --features metal,accelerate models::gpt2 (25 passed)
• DEVELOPER_DIR=... cargo clippy --release --lib --tests --features metal,accelerate (clean; the two host_preprocessor warnings are pre-existing on main)
• cargo fmt --all -- --check
• python3 scripts/ci/check_cross_repo_refs.py

Closes #925

[inureyes]

One acceptance criterion is intentionally left unticked on the issue: the end-to-end run that loads a GPT-2 checkpoint with an adapter through load_model_with_adapter and generates against it. No GPT-2 checkpoint is present in this environment, and real-model validation is handled separately from this change.

What that run should show, given the guard: a raw HuggingFace gpt2 checkpoint plus an adapter targeting any of h.N.attn.c_attn, h.N.attn.c_proj, h.N.mlp.c_fc or h.N.mlp.c_proj now fails at load with the "LoRA fusion refused" error naming h.0.attn.c_attn.weight, its [n_embd, 3 * n_embd] shape, and the affected targets, instead of aborting the process (non-square) or loading a silently corrupted attn.c_proj (square). An mlx-community GPT-2 conversion plus the same adapter fuses and generates exactly as before, since its projections are already [out_features, in_features] and the layout probe finds no Conv1D signal.

[inureyes]

CI blocker, pre-existing on main, unrelated to this PR

Pipeline Parallel CI / 2-host logical is red. The pipeline-parallel job itself passed; only the trailing Clippy on PP modules step failed, on two findings in src/multimodal/:

error: unused import: `export_qwen2_vl_prefill`
  --> src/multimodal/host_preprocessor.rs:47:51
error: function `export_qwen2_vl_prefill` is never used
  --> src/multimodal/host_preprocessor_export.rs:154:15

This PR touches five files, none of them under src/multimodal/:

src/distributed/pipeline/partial_loading_adapter_tests.rs
src/lora/loader.rs
src/lora/loader_tests.rs
src/models/gpt2.rs
src/models/gpt2_tests.rs

The cause is on main. export_qwen2_vl_prefill has exactly one non-test call site, src/multimodal/host_preprocessor.rs:428, and it sits inside a block gated on #[cfg(feature = "xla-iree")]. The CI step runs cargo clippy -p mlxcel --lib --tests -- -D warnings with no features, so that call site is compiled out, the import goes unused and the function goes dead, and -D warnings turns both into errors. It reproduces locally on the same command with no local changes applied.

It has gone unnoticed because this workflow is path-filtered to the PP modules and last ran on 2026-07-20, while the code arrived on 2026-07-24 in edc0ebbb6 (feat: add Qwen2-VL OpenXLA vision path). This is simply the first PP-touching PR since then.

Left unfixed here rather than folded in, since it is an unrelated module and belongs in its own change. The fix is to gate the import and the function on feature = "xla-iree" the way the call site already is, and it will unblock every PP-touching PR, not just this one.