Skip to content

fix(rt-detr-v2): needs_sanitize returns on the first marker found in HashMap order, so its verdict is order-dependent #1276

Description

@inureyes

Problem / Background

needs_sanitize (src/vision/detection/rt_detr_v2/sanitize.rs:46-73) decides whether a freshly loaded weight map is raw HuggingFace layout and therefore needs the rename plus transpose pipeline. It walks weights.keys(), and WeightMap is a std::collections::HashMap (src/lib/mlxcel-core/src/weights.rs:29: pub type WeightMap = HashMap<String, UniquePtr<MlxArray>>), so the walk order is randomized by RandomState, per map instance rather than merely per process. Inside the loop the function sets has_mlx_marker for keys starting with vision.backbone. or vision.hybrid_encoder. (lines 50-52), sets has_hf_marker for six HF markers (lines 53-61), and then returns immediately: if has_mlx_marker { return false } at lines 63-65, followed by if has_hf_marker { return true } at lines 66-68.

The comment at line 62 says "Early out once both decisions are unambiguous", but the code does not wait for both. It returns as soon as EITHER flag is set. For any weight map that contains both marker families, the answer is decided by whichever key the randomized walk happens to reach first, so the same file can produce different verdicts on different runs. The comment should be treated as part of the defect, since it describes a guard the code does not implement.

A wrong answer is expensive in one direction. The module documentation at src/vision/detection/rt_detr_v2/sanitize.rs:30-32 states that MLX-layout checkpoints must skip the pipeline because "re-running it would double-transpose conv weights". A spurious true therefore runs the rename and transpose over already-transposed weights and silently corrupts them: the model still loads through RtDetrV2Model::load (src/vision/detection/rt_detr_v2/model.rs:110-115) and still emits detections, they are just wrong. A spurious false on a genuine HF checkpoint is the milder direction: the rename never runs and the load fails or mismatches loudly.

Reachability, stated precisely so this is not overclaimed. Two of the HF markers are unanchored substring tests, key.contains(".convolution.") (line 55) and key.contains(".normalization.") (line 56), so they can match anywhere in a key rather than only at a prefix. A correctly converted MLX checkpoint should not trip them, because rename_stripped in the same file (lines 102-103, reached through the public rename_key at lines 163-165) rewrites .convolution. to .conv. and .normalization. to .bn.. So for well formed inputs on both sides this is latent rather than live. It becomes reachable when a weight map carries both families at once: a partially converted or hand merged checkpoint, an MLX checkpoint shipping extra auxiliary tensors under HF field names, or a future MLX layout that keeps any .normalization. shaped key. The point of this issue is that the function currently has no defense against that case and fails toward the corrupting direction at random rather than deterministically or loudly.

One more detail worth fixing while there: the trailing has_hf_marker return at line 72 can only be reached when no key set either flag, in which case has_hf_marker is necessarily false. It reads as though it could return true and cannot. Returning false explicitly, keeping the existing "treat as already-MLX to avoid corrupting an unexpected layout" reasoning from the comment at lines 70-71, would say what actually happens.

Current Behavior

pub fn needs_sanitize(weights: &WeightMap) -> bool {
    let mut has_mlx_marker = false;
    let mut has_hf_marker = false;
    for key in weights.keys() {
        if key.starts_with("vision.backbone.") || key.starts_with("vision.hybrid_encoder.") {
            has_mlx_marker = true;
        }
        if key.starts_with("model.backbone.")
            || key.starts_with("backbone.model.")
            || key.contains(".convolution.")
            || key.contains(".normalization.")
            || key.starts_with("model.decoder.")
            || key.starts_with("model.encoder.")
        {
            has_hf_marker = true;
        }
        // Early out once both decisions are unambiguous.
        if has_mlx_marker {
            return false;
        }
        if has_hf_marker {
            return true;
        }
    }
    // No markers at all: treat as already-MLX (no-op) to avoid corrupting an
    // unexpected layout.
    has_hf_marker
}

The only caller is RtDetrV2Model::load (src/vision/detection/rt_detr_v2/model.rs:110-115), which runs sanitize::sanitize(raw) when the verdict is true and passes raw through untouched when it is false. The existing coverage, needs_sanitize_detects_layout (src/vision/detection/rt_detr_v2/sanitize.rs:302-321), builds one single-key MLX map and one single-key HF map, so it never exercises a map holding both families and cannot observe the ordering dependence.

Proposed Solution

Maintainer's choice on the details, but the shape should be: scan every key before deciding rather than returning from inside the loop, then apply an explicit precedence so the verdict is independent of iteration order. MLX should win, since that is the safe direction: skipping the pipeline on a genuine HF checkpoint fails loudly at weight lookup, while running it on an MLX checkpoint corrupts conv weights silently.

Decide separately whether a map carrying both families should be rejected loudly instead of resolved silently. That input is a checkpoint the loader does not actually understand, and needs_sanitize returns a bare bool today, so surfacing it would mean changing the signature (for example to Result<bool, String>, matching the Result<Self, String> error style already used by RtDetrV2Model::load) and threading the error out through the caller. If the decision is to keep the silent resolution, say so in the code comment rather than leaving the current misleading "both decisions" wording.

Also replace the trailing has_hf_marker at line 72 with an explicit false, and correct the line 62 comment so it describes what the code does.

Scope

In scope: src/vision/detection/rt_detr_v2/sanitize.rs (the needs_sanitize body, its doc comment, the two inline comments at lines 62 and 70-71), the inline mod tests in the same file (new regression test), and src/vision/detection/rt_detr_v2/model.rs:110-115 only if the signature changes to a Result.

Out of scope: the rename rules in rename_stripped, the transpose in maybe_transpose_conv, the drop rules in should_drop, and any change to the marker set itself. Widening or anchoring the six HF markers is a separate design question and should not ride along here. The sibling HashMap ordering instances tracked in #1267 and in the nodes_with_role failover ordering issue (#1277 (src/distributed/registry.rs:146)) filed alongside this one as #1277 are also out of scope.

Implementation Notes

  • Reuse: keep the existing marker string set exactly as it is, and keep the existing "treat as already-MLX to avoid corrupting an unexpected layout" reasoning for the no-marker case. This is a control-flow fix, not a re-derivation of the layout sniffing rules.
  • Constraints: the verdict must not depend on weights.keys() order for any input. Do not switch WeightMap to an ordered map to work around this: it is a shared type used across the whole loader (src/lib/mlxcel-core/src/weights.rs:29) and changing it here would be a cross-cutting change well outside this issue.
  • Edge cases: empty map (no keys) keeps the current verdict, false; map with only MLX markers is false; map with only HF markers is true; map with neither is false; map with both is the case this issue exists for, and must be either a stable false or a loud error, never a coin flip.
  • Error handling: if the both-families case is made loud, the message must name at least one offending key from each family so the operator can see which tensors caused the ambiguity, and RtDetrV2Model::load must propagate it rather than falling back to a default verdict. If it stays silent, no behavior change is needed at the call site.

Acceptance Criteria

  • needs_sanitize inspects all keys before returning, so its verdict is a pure function of the key set and not of HashMap iteration order.
  • A weight map containing both an MLX-prefixed key (for example vision.backbone.embedder.embedder.0.conv.weight) and a .convolution. shaped key yields one stable verdict, or a deterministic error, across repeated construction of the map.
  • The line 62 comment no longer claims a "both decisions" guard that the code does not implement.
  • The trailing return at line 72 is an explicit false rather than a variable that is provably false at that point.
  • A regression test in the inline mod tests of sanitize.rs builds a mixed-marker map and asserts the same verdict over many repetitions (rebuilding the map each iteration, since RandomState varies per map instance, so a single call cannot distinguish a fixed answer from a lucky one).
  • The existing needs_sanitize_detects_layout test still passes unchanged.
  • Integrated into the real load path (signature kept as bool, so RtDetrV2Model::load is untouched and compiles clean; not exercised against actual checkpoint files, see comment below): RtDetrV2Model::load still compiles and behaves correctly against both a pre-converted MLX checkpoint and a raw HF checkpoint (signature change threaded through if the loud-error option is taken).

Verification

# Targeted unit tests (macOS)
cargo test --features metal,accelerate rt_detr_v2::sanitize

# Targeted unit tests (Linux / CUDA)
cargo test --features cuda rt_detr_v2::sanitize

# Full local gate
make verify            # macOS: verify-versions, verify-kernel-dtype-keys, verify-fmt, verify-clippy, verify-test
make verify-test-cuda  # Linux/NVIDIA merge gate

A pass means the targeted filter runs the existing needs_sanitize_detects_layout plus the new mixed-marker regression test with zero failures, and the repetition loop in the new test never observes a differing verdict. Repeating the targeted command several times in a row is a useful sanity check, since the defect this fixes is exactly the kind that a single green run cannot rule out.

Technical Considerations

Found during the review of PR #1268 while sweeping for the root-cause class fixed by #1265 and PR #1266 (HashMap iteration order leaking into an ordered or order-sensitive decision). Sibling instances: #1267 in src/lang_bias.rs, and the nodes_with_role failover ordering issue in #1277 (src/distributed/registry.rs:146) filed alongside this one as #1277.

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:modelsModel architectures, weights, loading, metadatapriority:mediumMedium prioritystatus:doneCompletedtype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions