fix(lora): fail the load when staged runtime adapter terms are never claimed
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Summary
Issue #1328 and PR #1576 made both LoRA paths refuse adapter tensors that do not map onto a base weight, instead of warn-skipping them and serving the base model under the adapter's name. The unfused runtime path has a second, independent way to reach the same outcome that #1576 does not close: a term that stages successfully but that no layer constructor ever claims is reported as a `warn!` and the load succeeds anyway. This issue owns that hole, plus two smaller adapter-validation findings from the same review.
Line numbers below are against PR #1576's branch (`fix/issue-1328-lora-adapter-validation`) and will shift when it squash-merges. The function names are the durable anchors. `src/loading/mod.rs` and `src/lora/loader.rs` are cited by function name only because that branch still has pending edits in both.
## Current behavior
`stage_runtime_adapters` in `src/lora/runtime.rs:286` validates every adapter tensor against the base weight map, then stages per-layer terms into a thread-local keyed by base weight prefix (the base weight name minus its `.weight` suffix) through `mlxcel_core::runtime_lora::stage` (`src/lib/mlxcel-core/src/runtime_lora.rs:183`). Staging is only half of the contract: a staged term does nothing unless a layer constructor calls `mlxcel_core::runtime_lora::claim` (`src/lib/mlxcel-core/src/runtime_lora.rs:200`) with that exact prefix during model construction.
Exactly three constructors claim, all in `src/lib/mlxcel-core/src/layers.rs`: `Linear::from_weights` (`layers.rs:1031`), `UnifiedLinear::from_weights` (`layers.rs:2016`), and `FusedQKVLinear` (`layers.rs:2879-2881`, which claims the three separate q/k/v prefixes because the fused representation concatenated the planes). Any projection a family builds outside those three leaves its term unclaimed. `SwitchLinear::from_weights` (`src/models/switch_layers.rs:382`), which is how every MoE family loads its expert weights, does not claim; neither does any hand-rolled `from_weights` that pulls the tensor out of the `WeightMap` itself.
`mlxcel_core::runtime_lora::drain_unclaimed` (`src/lib/mlxcel-core/src/runtime_lora.rs:222`) returns those orphaned prefixes. `finish_runtime_staging` (`src/lora/runtime.rs:352`) iterates them and emits one `tracing::warn!` each, returning `()`. `load_model_with_adapter_specs` in `src/loading/mod.rs` calls it for its side effect and discards nothing because there is nothing to discard. The net result is the same failure mode #1328 was filed about: the adapter validates, the log line says `Staged runtime LoRA adapter for N layers`, and those N layers serve base weights at every forward.
This is not an edge case. `runtime_lora_mode` in `src/server/cli_input.rs:664` selects the unfused runtime path whenever the b10621 `--lora` or `--lora-scaled` spellings are used without `--adapter`, without `--lora-fuse`, and without tensor or pipeline parallelism, which is the default for the b10621 CLI surface (`src/lora/multi.rs` parses those spellings). An operator who starts `mlxcel-server --lora ` against a family whose targeted modules are built outside the three claiming constructors gets a server that answers from the base model and logs a warning about it.
Second finding, smaller. The fused single-process path and the pipeline-parallel stage path validate adapters against different weight maps. `load_model_with_adapter` and `load_model_with_adapter_specs` in `src/loading/mod.rs` both read the raw `mlxcel_core::weights::load_weights_from_dir` and only reach `models::sanitize_tied_embeddings` (`src/models/sanitize.rs:1502`) later, inside `load_model_from_weights`. The pipeline stage path loads through `models::load_text_weights` (`src/distributed/pipeline/stage_executor/llama.rs:48`), which calls `sanitize_tied_embeddings` at `src/models/sanitize.rs:1644`, before the adapter is composed. On a `tie_word_embeddings` checkpoint the sanitizer synthesizes `lm_head.weight` from `model.embed_tokens.weight`, so an adapter carrying an `lm_head` pair now resolves under pipeline parallelism and hard-fails under single-process. Same adapter, same checkpoint, two verdicts.
Third finding, recorded rather than fixed here. MoE LoRA is unsupported on both paths. `mlx-lm`'s `LoRASwitchLinear` writes 3-D `lora_a` and `lora_b` tensors, one slice per expert, and `compute_lora_delta` in `src/lora/loader.rs` bails with `Expected 2D LoRA matrices, got lora_a={...}, lora_b={...}` for anything that is not rank 2. That was already a hard error before #1576, so it is not a regression introduced by that work, but it is undocumented and an operator meets it as a raw shape complaint rather than as a stated limitation.
## Expected behavior
An unclaimed staged term fails the load. `finish_runtime_staging` returns the unclaimed prefixes and `load_model_with_adapter_specs` refuses the load when the list is non-empty, naming every unclaimed layer in one error, in the same posture #1576 established for unmapped tensors: report all offenders at once rather than one per load attempt, because an adapter built against the wrong module set is wrong on every layer it carries.
The fused single-process path and the pipeline stage path agree on whether an `lm_head` pair is applicable to a tied-embeddings checkpoint. Whichever answer is chosen, both paths give it.
The MoE limitation is stated in the user-facing LoRA documentation, and the error an operator sees names it as unsupported rather than reporting a tensor rank.
## Implementation notes
- **Return the list, do not log it.** Change `finish_runtime_staging` to `#[must_use] pub fn finish_runtime_staging() -> Vec` and sort the result before returning: `drain_unclaimed` drains a `HashMap`, so the order is nondeterministic and an unsorted error message would be flaky across runs and untestable with more than one unclaimed layer. `src/lora/runtime_tests.rs::a_well_formed_runtime_adapter_still_stages` already sorts for exactly this reason and is the foothold to extend.
- **Keep the drain on the error path.** The current call site in `load_model_with_adapter_specs` deliberately calls `finish_runtime_staging` before propagating a construction failure, so a failed load cannot leak terms into the next one on the same thread. That ordering must survive the change: bind the returned list, keep the existing `model?`, then fail on a non-empty list.
- **Error shape.** Mirror `validate_adapter_tensors`: one error listing every offender, one per line, naming the layer prefix and stating that the model builds that layer outside the claiming constructors. The existing `warn!` text in `src/lora/runtime.rs:354-358` already says this and can be lifted into the error.
- **Why #1576 did not do it.** #1576 was scoped to tensor-to-base-weight mapping, which is checkable against the weight map alone. This one is a hard-failure change to the default server adapter channel, gated on which constructors a family happens to use, so it cannot be validated without real checkpoints across several families. A family that legitimately builds a targeted module outside the three constructors would start refusing an adapter that used to load, and the fix in that case is to make the constructor claim, not to relax the check.
- **Reuse.** Do not add a second unclaimed-reporting path. `drain_unclaimed` stays the single source; only its consumer changes.
- **Tied-embeddings decision needs thought, not a call reorder.** Neither direction is obviously right. Applying an `lm_head` delta to a tied checkpoint breaks the tie between the embedding and the output projection, which is a silent numerical change, so making single-process match the pipeline path is not automatically the correct fix. Refusing on both paths is defensible. Note that stock `mlx_lm.lora` only wraps modules inside `model.layers`, so no adapter it produces carries an `lm_head` pair and the divergence is latent today; that is what makes it safe to decide deliberately rather than urgently.
- **Verification.**
```
cargo test --profile test-fast --features metal,accelerate --lib lora
cargo test --profile test-fast --features metal,accelerate --lib loading::
cargo clippy --workspace --all-targets --features metal,accelerate -- -D warnings
cargo fmt --all -- --check
```
- **Real-checkpoint validation is required and cannot be skipped.** Start `mlxcel-server --lora ` against at least one dense family whose targeted modules go through the claiming constructors (must still load, and must produce output that differs from the same prompt with no adapter) and at least one family that builds a targeted module elsewhere (must now fail with the layer named). A unit test with a synthetic `WeightMap` cannot cover this, because what a family's constructors claim is a property of the real model code.
## Acceptance criteria
- [ ] `finish_runtime_staging` returns the unclaimed layer prefixes, is marked `#[must_use]`, and sorts the list before returning.
- [ ] `load_model_with_adapter_specs` fails the load when that list is non-empty, with one error naming every unclaimed prefix, and still drains staged terms on the model-construction error path.
- [ ] A unit test stages a term under a prefix no constructor claims and asserts the load fails with that prefix named; another asserts several unclaimed prefixes are reported together in sorted order.
- [ ] A unit test asserts the positive control still passes: an adapter whose every prefix is claimed loads and leaves the unclaimed list empty.
- [ ] Starting `mlxcel-server --lora ` on a real checkpoint whose targeted modules are all claimed still loads and changes generated output versus no adapter, verified by hand on the built binary.
- [ ] Starting `mlxcel-server --lora ` on a real checkpoint where a targeted module is built outside the claiming constructors fails at load with that layer named, verified by hand on the built binary.
- [ ] The `lm_head`-under-tied-embeddings divergence is resolved to one verdict, both `src/loading/mod.rs` and `src/distributed/pipeline/stage_executor/llama.rs` reach it, and a test pins the agreed answer on both paths.
- [ ] The user-facing LoRA documentation states that MoE expert adapters (`mlx-lm` `LoRASwitchLinear`, 3-D `lora_a` / `lora_b`) are unsupported, and the error an operator hits names the limitation instead of reporting a tensor rank.
- [ ] `cargo clippy --workspace --all-targets --features metal,accelerate -- -D warnings` and `cargo fmt --all -- --check` are clean.
## Out of scope
- DoRA fusion. #1576 refuses DoRA by name on both paths; actually supporting the magnitude vectors is separate work.
- Fused adapter application on a quantized base. It fails on the pre-existing shape guard because a 4-bit `.weight` is packed (`[out, in/8]` as `U32`) while the delta is dense `[out, in]`, which predates #1576 and is not what this issue is about. The unfused runtime channel is the supported route for a quantized base.
- Model-local fused kernels that read a `UnifiedLinear`'s quantized parts directly and never consult `mlxcel_core::runtime_lora::any_active`, such as the Mamba-2 fused path in `src/models/nemotron_h.rs`. Those terms are claimed, so an empty unclaimed list does not by itself prove the adapter reaches the forward pass. That is a distinct hole and needs its own issue.
Contributor guide
Research direction
Start with finish_runtime_staging in src/lora/runtime.rs and its call in load_model_with_adapter_specs in src/loading/mod.rs, then read the claiming constructors and runtime_lora tests. Run the lora and loading test commands listed in the issue, adding coverage for sorted unclaimed prefixes and the positive control. Done means unclaimed terms fail together, tied-embedding behavior agrees across paths, and real dense and MoE checkpoints verify the documented LoRA behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 45/100