Skip to content

feat: add shared DINOv2 feature encoder - #2549

Merged
inureyes merged 5 commits into
mainfrom
feature/issue-2350-dinov2-backbone
Oct 22, 2026
Merged

inureyes merged 5 commits into
mainfrom
feature/issue-2350-dinov2-backbone

Conversation

@inureyes

@inureyes inureyes commented Oct 22, 2026 •

Copy link
Copy Markdown
Member

Hugging Face DINOv2 checkpoints previously had no shared encoder or standalone runtime in mlxcel, which blocked the dense-model families that depend on the backbone. This adds a reusable standard/register-token encoder and exposes image feature extraction through mlxcel predict.

The encoder loads public Hugging Face standard and register checkpoints, remaps and validates their weights, supports boolean patch masks, preserves register-token ordering, returns normalized intermediate layers, and implements rectangular positional interpolation in both upstream modes. Shared consumers can construct a backbone from a scoped parent checkpoint without consuming the parent weight map. The giant preset uses its SwiGLU feed-forward layout.

The standalone predictor reads preprocessor_config.json and reproduces the slow Hugging Face BitImageProcessor path: RGB conversion, shortest-edge resize, Pillow-compatible bicubic uint8 rounding, center crop, rescaling, and normalization. It writes last_hidden_state, pooler_output, and spatial-only hidden_patch_tokens. Architecture and operator documentation describe both supported checkpoint types.

Pinned full-tensor CPU f32 fixtures cover facebook/dinov2-small and facebook/dinov2-with-registers-small after real checkpoint preprocessing. Tests also cover predictor preprocessing, positional interpolation literals, [4, 11] intermediate outputs, register exclusion, f16 masking, scoped loading, allocation limits, invalid public-layer inputs, and sanitizer behavior. Checkpoint-driven image, token, architecture, and intermediate-output sizes are bounded before expensive allocations; attention uses the shared safe dispatch.

Validation:

  • cargo test -p mlxcel --profile test-fast --features rocm --lib dinov2 -- --nocapture --test-threads=1: 13 passed, 0 failed, 0 ignored. The explicit CPU f32 reference tests produced relative RMS values from 1.197926287e-7 for preprocessing through 4.275870861e-6 for model outputs.
  • cargo clippy -p mlxcel --profile test-fast --features rocm --lib --tests -- -D warnings
  • cargo build -p mlxcel --profile test-fast --features rocm --bin mlxcel
  • Guarded MLXCEL_DEVICE=gpu CLI predictions on gfx1151 for both checkpoints wrote exactly meta.json and result.safetensors. Standard-checkpoint tensor RMS values were 2.844445329e-6, 2.473357604e-6, and 2.845811506e-6; register-checkpoint values were 4.548628530e-6, 2.131596900e-6, and 4.596212478e-6 for last hidden state, pooler output, and patch tokens respectively.
  • mlxcel arch --json lists dinov2 with both model types and the predict/image/embeddings capability.
  • cargo fmt --all -- --check
  • python3 scripts/ci/check_env_var_docs.py
  • python3 scripts/ci/check_binary_assets.py
  • git diff --check
  • Independent implementation and security/performance reviews completed with all findings resolved.

Closes #2350

@inureyes inureyes added status:in-progress Currently being worked on type:enhancement New features, capabilities, or significant additions priority:low Low priority area:models Model architectures, weights, loading, metadata modelsize:small 4-bit checkpoint <= 10GB; fast iteration, safe for smoke tests modeltype:vlm Vision-language model labels Oct 22, 2026
@inureyes
inureyes marked this pull request as ready for review October 22, 2026 18:18
@inureyes inureyes added status:done Completed and removed status:in-progress Currently being worked on labels Oct 22, 2026
inureyes added a commit that referenced this pull request Oct 22, 2026
Port standard and register-token Hugging Face DINOv2 checkpoints with shared dense-model APIs, upstream-compatible interpolation, masking, intermediate features, and giant SwiGLU support.

Expose standalone image feature extraction through mlxcel predict, register architecture capabilities, and add pinned full-output reference fixtures plus operator documentation.
Load and enforce the checkpoint image processor policy, reproduce Pillow's bicubic uint8 rounding, and cover predictor output against real references. Bound checkpoint-driven allocations and harden the shared backbone APIs.
Describe the checkpoint processor policy and its Hugging Face defaults instead of the superseded patch-alignment behavior.
@inureyes
inureyes force-pushed the feature/issue-2350-dinov2-backbone branch from f084c05 to ab61c8f Compare October 22, 2026 18:28
@inureyes
inureyes merged commit c6ca6c4 into main Oct 22, 2026
29 checks passed
@inureyes
inureyes deleted the feature/issue-2350-dinov2-backbone branch October 22, 2026 18:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:models Model architectures, weights, loading, metadata modelsize:small 4-bit checkpoint <= 10GB; fast iteration, safe for smoke tests modeltype:vlm Vision-language model priority:low Low priority status:done Completed type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: Port DINOv2 backbone (dinov2, dinov2_with_registers) for dense models

1 participant