lablup / lablup/mlxcel

refactor(embeddings): hoist family-local weight-key handling into the shared embedding sanitizer

Open
#1,424 0 comments 0 reactions 0 assignees View on GitHub
area:models priority:low status:ready type:refactor
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Problem / Background
The epic #1348 embedding wave landed several families in parallel, and to avoid merge conflicts each PR kept some shared-looking logic local to its own module. The result is three places that assemble a headless Llama backbone, a weight-key rule that lives outside the shared sanitizer, and a `max_length` derivation that ignores the one file the Llama-Nemotron reference reads. None of this changes model output today, but every future family will copy one of the three variants, and a bug fix in one will not reach the others.

## Current Behavior
- Shared key normalization: `src/models/embedding_sanitize.rs:41` (`BACKBONE_ROOTS = ["embed_tokens.", "layers.", "norm."]`), `:44` (`HEAD_ROOTS`), `:105-121` (`prefix_backbone_roots`), `:123-134` (`sanitize_decoder_embedding_weights`, which folds `Dense` folders, drops the generation head and prefixes bare roots with `model.`).
- LFM2: `src/models/lfm2_embedding.rs:68-72` declares `LFM2_EXTRA_BACKBONE_ROOT = "embedding_norm."` and `:87-99` (`prefix_embedding_norm`) re-implements the prefixing loop for that one root, called at `:149` after the shared sanitize at `:147`. PR #1415 kept this local to avoid conflicting with the parallel wave.
- Llama assembly, three ways: `src/models/llama_bidirec.rs:70-77` (`LlamaBidirecModel`: `embed_tokens`, `layers: Vec`, `norm`) built at `:171-215` from `UnifiedEmbedding::from_weights`, `TransformerBlock::from_weights_with_rope`, and `model.norm.weight`; `src/models/headless_llama.rs:35-39` (`HeadlessLlama`, PR #1414, used by ColIdefics3) built at `:43-71` from the exact same three calls; and `src/models/llama_nemotron_vl_embedding.rs:83-95` (`LlamaNemotronVLEmbeddingModel`) which instead holds `text: Llama3Model` (built at `:229` with `Llama3Model::from_weights`, so it also loads the `lm_head` that the embedder never uses) and drives `self.text.layers` and `self.text.norm` by hand at `:258-270` (`forward_text`).
- Key sanitizers, two ways: `src/models/llama_bidirec.rs:66-67` (`LANGUAGE_MODEL_PREFIX`, `DROPPED_BUFFERS = ["rotary_emb.inv_freq", "position_ids"]`) with `sanitize_llama_bidirec_weights` at `:89-106`, and `src/models/llama_nemotron_vl_embedding.rs:110-131` (`sanitize_nemotron_vl_weights`) which strips the same `language_model.` prefix and drops the same two buffers with its own loop before calling the shared sanitize.
- `max_length`: `src/embeddings/loader.rs:239-250` calls `EmbeddingLimits::derive` (`src/embeddings/limits.rs:52-66`), which calls `derive_max_length` (`limits.rs:86-125`): the minimum of `EMBEDDING_MAX_LENGTH_CAP` (8192, `limits.rs:29`), `sentence_bert_config.json` `max_seq_length`, `tokenizer_config.json` `model_max_length`, `config.json` `max_position_embeddings` for absolute-position families, and the operator override. `processor_config.json` is not consulted. The Llama-Nemotron-VL family reads that file only for tiling and prompt keys (`src/models/llama_nemotron_vl_embedding.rs:135-166`, `read_processor_config`), so the family runs at 8192 while the reference processor caps passages and queries lower (the reported reference values are 4096 for passages and 512 for queries; confirm the exact key names in the checkpoint's `processor_config.json` before implementing). The `EmbeddingModel::max_sequence_length` hook (`src/embeddings/model.rs:128`, consumed at `loader.rs:247-249`) already lets a family lower the derived limit.

## Proposed Solution
1. Key normalization: extend `BACKBONE_ROOTS` in `src/models/embedding_sanitize.rs` with `"embedding_norm."`, delete `LFM2_EXTRA_BACKBONE_ROOT` and `prefix_embedding_norm` from `src/models/lfm2_embedding.rs`, and cover the root in `embedding_sanitize`'s tests. Move `LANGUAGE_MODEL_PREFIX` and `DROPPED_BUFFERS` into `embedding_sanitize.rs` as `pub(crate) fn strip_language_model_wrapper(&mut WeightMap)` and `pub(crate) fn drop_derived_buffers(&mut WeightMap)`, and have both `sanitize_llama_bidirec_weights` and `sanitize_nemotron_vl_weights` call them in the documented order (wrapper strip, buffer drop, shared sanitize). Idempotency on an already prefixed map is part of the contract and must keep its test.
2. Headless Llama: make `HeadlessLlama` (`src/models/headless_llama.rs`) the single assembly. `LlamaBidirecModel` holds `backbone: HeadlessLlama` plus `pooling`, `normalize`, `embedding_dim`, and its `forward_hidden` delegates to `HeadlessLlama::forward_hidden` with the bidirectional mask. `LlamaNemotronVLEmbeddingModel` replaces `text: Llama3Model` with `text: HeadlessLlama`, uses `HeadlessLlama::embed_tokens` where it currently calls `self.text.embed_tokens.forward` and `HeadlessLlama::forward_hidden` in place of `forward_text`. This drops the unused `lm_head` load for the VL family. `HeadlessLlama::forward_hidden` (`headless_llama.rs:83-89`) already takes `input_embeddings: Option<&MlxArray>` and `mask: Option<&MlxArray>`, so the bidirectional families pass `create_bidirectional_padding_mask` through it and the VL family passes its merged vision embeddings; no signature change is needed.
3. `max_length`: teach `derive_max_length` in `src/embeddings/limits.rs` to also read `processor_config.json` (`max_length` and, if present, the query-specific key) as another candidate in the minimum, behind the same `> 0` filter as the other sources, and add the file to the doc comment at `limits.rs:82-85`. If the reference distinguishes passage and query lengths, expose the query length through `EmbeddingLimits` as `query_max_length: Option` and have the engine apply it when `EmbedOptions` marks the input as a query; if it does not, a single `max_length` candidate is enough. State which one the checkpoint actually needs in the PR.

Rejected alternative: leaving `LlamaNemotronVLEmbeddingModel` on `Llama3Model` and only deduplicating the sanitizers. The by-hand layer loop at `llama_nemotron_vl_embedding.rs:266-268` is the third copy of the forward pass and is the part most likely to drift (mask handling, cache creation), so it is the one worth removing.

## Scope
**In scope:** `src/models/embedding_sanitize.rs` (and tests), `src/models/lfm2_embedding.rs`, `src/models/llama_bidirec.rs`, `src/models/headless_llama.rs`, `src/models/llama_nemotron_vl_embedding.rs`, `src/embeddings/limits.rs` (and tests), `src/embeddings/loader.rs` if `EmbeddingLimits` grows a field, `src/embeddings/engine.rs` only if a query-specific limit is needed.

**Out of scope:** ColIdefics3's use of `HeadlessLlama` (already the target shape), the generation-side `Llama3Model`, any change to pooling, normalization or prompt formatting.

## Implementation Notes
- **Reuse**: `HeadlessLlama::from_weights` / `make_caches` / `embed_tokens` / `forward_hidden` (`headless_llama.rs:44-89`); `create_bidirectional_padding_mask` already used by both bidirectional families; `sanitize_decoder_embedding_weights` as the single entry point after the wrapper strip.
- **Constraints**: bit-identical output for the three real checkpoints (see acceptance); the `embedding_norm.` root must be prefixed before `Lfm2Model::from_weights` runs its own sanitize (`lfm2_embedding.rs:135-138`), so ordering inside the shared helper must be preserved; `Llama3Model` stays untouched.
- **Edge cases**: a map that already carries `model.` on every key (idempotent, zero renames); a checkpoint with a `language_model.model.` double wrapper (strip exactly the documented prefix, do not loop); `processor_config.json` absent or without a length key (fall through to the existing candidates); a `processor_config.json` value larger than the tokenizer's (the minimum still wins).
- **Error handling**: unchanged; the loaders keep returning `anyhow` errors naming the missing key, and the sanitizer never panics on an unexpected key.

## Acceptance Criteria
- [ ] `grep -rn "embedding_norm\|inv_freq\|language_model\." src/models/lfm2_embedding.rs src/models/llama_bidirec.rs src/models/llama_nemotron_vl_embedding.rs` returns only doc comments; the constants and loops live in `src/models/embedding_sanitize.rs`.
- [ ] `LlamaBidirecModel` and `LlamaNemotronVLEmbeddingModel` both hold a `HeadlessLlama` and neither constructs `TransformerBlock` or `Llama3Model` directly.
- [ ] `derive_max_length` reads `processor_config.json`; a unit test next to `derive_max_length` (there is no `src/embeddings/limits_tests.rs` today; create it with the crate's `#[path]` sibling-file convention) with a temp dir containing only `processor_config.json` asserts the derived value equals that file's limit.
- [ ] The real-checkpoint gates `models::lfm2_embedding_tests::lfm2_embedding_checkpoint_ranks_the_related_passage_first` (`src/models/lfm2_embedding_tests.rs:424`), `models::llama_bidirec_tests::llama_nemotron_embed_checkpoint_ranks_the_related_passage_first` (`src/models/llama_bidirec_tests.rs:360`), and `models::llama_nemotron_vl_embedding_tests::llama_nemotron_vl_text_gates_hold_on_the_real_checkpoint` plus `..._image_gates_hold_on_the_real_checkpoint` (`src/models/llama_nemotron_vl_embedding_tests.rs:358` and `:418`) pass, and a before/after comparison of the emitted vectors for the same inputs (dump to a file on `main`, dump on the branch, compare) is bit-identical for LFM2.5-Embedding, llama-nemotron-embed-1b-v2 and llama-nemotron-embed-vl-1b-v2, except where the new `max_length` truncates an input that was previously longer than the reference limit (report those inputs explicitly).
- [ ] Integrated into the real code flow: `mlxcel embed -m ` reports the new `max_length` and `POST /v1/embeddings` truncates at it.

## Verification
```bash
cargo fmt --all -- --check
cargo clippy --lib --tests --profile test-fast --features cuda -- -D warnings
cargo test --profile test-fast --features cuda --lib models::embedding_sanitize
cargo test --profile test-fast --features cuda --lib embeddings::limits
cargo test --profile test-fast --features cuda --lib -- --test-threads=1 models::lfm2_embedding models::llama_bidirec models::llama_nemotron_vl_embedding models::headless_llama models::colidefics3
```
Pass: every command exits 0 and the three real-checkpoint gates report the same scores as on `main` (the gates soft-skip when the checkpoints are absent, so run on the GB10 box where they are present).

## Technical Considerations
Related: PR #1414 (`HeadlessLlama`, ColIdefics3), PR #1415 (LFM2.5-Embedding and the local `embedding_norm.` rule), PR #1416 (Llama-Nemotron-VL-Embed), `src/models/embedding_sanitize.rs` as the shared home introduced earlier in the epic.

Originating PRs: #1414, #1415, #1416.

Follow-up from epic #1348.

Contributor guide

Open the contributing guide

Research direction

Start with src/models/embedding_sanitize.rs and src/models/headless_llama.rs, then compare the LFM2, Llama Bidirec, and Llama-Nemotron-VL call sites named in the issue. Read the limits implementation and its existing tests before checking the processor_config.json requirement. Done means the shared paths are used, the new limits test and targeted model tests pass, and the listed checkpoint outputs remain unchanged except for documented truncation.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
machine-learning
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.