lablup / lablup/mlxcel

feat(plamo2vl): port PLaMo 2.1 VL (SigLIP tiles + RMSNorm/GELU adapter + PLaMo 2 hybrid decoder)

Open
#1,380 0 comments 0 reactions 0 assignees View on GitHub
arch:hybrid area:cli area:docs area:inference area:models modelsize:small modeltype:vlm priority:low status:ready type:enhancement
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Summary

PLaMo 2.1 VL (`pfnet/plamo-2.1-2b-vl`, `pfnet/plamo-2.1-8b-vl`, `model_type: "plamo2vl"`) is a Japanese/English single-image VLM: a SigLIP-so400m (384px, patch 14) tower run over dynamically tiled crops, a three-stage RMSNorm/bias/Linear/GELU/Linear adapter, and the PLaMo 2 Mamba+attention decoder that mlxcel already ships for text. mlxcel rejects the checkpoint at detection. This change adds the family end to end: detection, the tiling processor, the adapter, the fixed Japanese instruction prompt with `` / `` framing, and three gaps in the existing PLaMo 2 text path that block embedding injection (no `embed_tokens` / `forward_with_embeddings`, a hardcoded RoPE base, and a tokenizer that cannot emit the image placeholder ids).

## Current behavior

- `src/models/detection.rs:298` has `"plamo2" => Ok(ModelType::Plamo2)` and no `plamo2vl` arm, so `get_model_type` returns `Unsupported model type: plamo2vl`.
- `src/models/plamo2.rs:1152-1272` implements `LanguageModel for Plamo2Model` with `forward`, `make_caches`, sequence-state methods and `trim_internal_caches`, but no `embed_tokens` and no `forward_with_embeddings`. The trait defaults (`src/lib/mlxcel-core/src/generate.rs:414-431`) ignore injected embeddings and return `None` from `embed_tokens`, so a VLM wrapper around this decoder would silently run text-only.
- `src/models/plamo2.rs:1024` builds every attention layer with `rope_base: 10000.0`. The VL text config carries `rope_theta: 1000000.0` and `rope_local_theta: 1000000.0` (both 2B and 8B); loading it through the current constructor would rotate at the wrong base and produce finite but wrong logits.
- `src/tokenizer/mod.rs:686-769` (`build_plamo_tokenizer`) registers exactly four added tokens (`<|plamo:unk|>` 0, `<|plamo:bos|>` 1, `<|plamo:eos|>` 2, `<|plamo:pad|>` 3). `tokenizer.jsonl` has 100000 rows; the VL checkpoint defines ids 100000..100015 (``, ``, ``, ``) only in `tokenizer_config.json` (`added_tokens_decoder` / `additional_special_tokens`). Today `` tokenizes as ordinary text.
- Nearest existing VLM shapes: the SigLIP tower `src/vision/encoders/siglip.rs:484-571` (`SigLipVisionModel::from_weights_with_quant_and_gelu`), the InternVL dynamic-tile processor `src/vision/processors/internvl.rs:45-215` (`InternVLProcessor::preprocess_with_tiles`), and the per-family runtime/prompt pair `src/vision/internvl.rs:143` + `src/multimodal/internvl_prompt.rs:58`.

## Expected behavior

`mlxcel generate -m models/plamo-2.1-2b-vl --image photo.jpg -p "この画像を説明してください。"` loads the checkpoint, tiles the image, runs the SigLIP tower once over all tiles, projects each of the 729 patch features per tile into the 2048-wide (2B) or 4096-wide (8B) text space, scatters them onto the `` positions of the templated prompt, and decodes with the PLaMo 2 hybrid cache. Same behavior through `/v1/chat/completions` with an `image_url` part.

### config.json

Top level: `model_type: "plamo2vl"`, `architectures: ["Plamo2VLForCausalLM"]`, `bos_token_id: 1`, `eos_token_id: 1` (EOS is `<|plamo:bos|>`, not id 2), `pad_token_id: 3`, `image_token_index` (absent; falls back to `vision_config.image_token_id`).

`text_config` (2B / 8B): `hidden_size` 2048 / 4096, `num_hidden_layers` 32, `num_attention_heads` 32, `num_key_value_heads` 4, `hidden_size_per_head` 128, `intermediate_size` 5632 / 16384, `mamba_num_heads` 64, `mamba_d_state` 64, `mamba_d_conv` 4, `mamba_step` 2, `mamba_enabled` true, `rms_norm_eps` 1e-6, `vocab_size` 100032, `rope_theta` 1e6, `rope_local_theta` 1e6, `full_attention_idx` (2B: `[0,2,...,30]`, 8B: `[]`), `attention_window_size` 32768 (not applied: plain `KVCache`), `tie_word_embeddings` absent (default true). No `eos_token_id` inside `text_config`: the top-level value must be used.

`vision_config` is flattened with a `vision_encoder_` prefix and an empty `model_type`: `vision_encoder_hidden_size` 1152, `vision_encoder_num_hidden_layers` 27, `vision_encoder_intermediate_size` 4304, `vision_encoder_num_attention_heads` 16, `vision_encoder_patch_size` 14, `vision_encoder_image_size` 384, `vision_encoder_num_channels` 3, `vision_encoder_layer_norm_eps` 1e-6, `vision_encoder_hidden_act` `"gelu_pytorch_tanh"`, plus unprefixed `image_token_id` 100002, `image_feature_size` 1152, `image_proj_hidden_size` 2048 (2B) / 4096 (8B).

`preprocessor_config.json`: `size {384,384}`, `image_mean [0.5,0.5,0.5]`, `image_std [0.5,0.5,0.5]`, `rescale_factor 1/255`, `patch_size 14`, `downsample_ratio 1.0`, `tile_context_length 3072`, `max_image_width/height 1e8`.

### Vision tower

Plain SigLIP: conv patch embed (14x14 stride 14, with bias), learned absolute position embedding over 27x27 = 729 positions (no interpolation; every tile is exactly 384x384), 27 pre-LN encoder layers (LayerNorm eps 1e-6, attention with q/k/v/out bias, MLP fc1/fc2 with `gelu_pytorch_tanh`), then `post_layernorm`. Output per tile: `[729, 1152]`. The attention-pooling `head.*` weights are unused and must be skipped.

### Adapter (`image_proj`)

For `x: [N*729, 1152]`:

```
x = rms_norm(x, weight = image_proj.norm0.weight + 1.0, eps = rms_norm_eps) # PLaMo offset norm, offset 1.0
x = x + image_proj.bias0._bias # [1152]
x = x @ image_proj.linear1.weight.T + image_proj.bias1._bias # [*, image_proj_hidden_size]
x = gelu_exact(x) # erf GELU, not tanh
x = x @ image_proj.linear2.weight.T + image_proj.bias2._bias # [*, hidden_size]
```

`linear1` and `linear2` have no `.bias`; the separate `_bias` tensors are the biases. Output `[N*729, hidden_size]` is scattered, in order, onto the positions where `input_ids == 100002`. The count must match exactly or loading of the prompt is an error.

### Image processor and prompt

1. `format_prompt(text, image_length)` builds the fixed instruction template (newline-joined, always applied, no chat roles):

```
以下はタスクを説明する指示で、文脈を説明した入力とペアになっています。要求を適切に補完するよう応答を書いてください。

### 指示:
{text}

### 応答:

```

with `` repeated `image_length` times before the text.

2. Tiling (`image_size = 384`, `num_image_token = (384/14)^2 = 729`, `max_tiles = 24`, `context_length = 3072`, `use_thumbnail = true`): candidate grids `(cols, rows)` with `1 <= cols*rows <= max_num`, sorted by `cols*rows`; pick the grid maximizing `min(ar, cols/rows) / max(ar, cols/rows)` (ties go to the candidate when `w*h > 0.5 * 384^2 * cols * rows`). Resize the image to `(cols*384, rows*384)` (bicubic), crop `cols*rows` tiles row-major, and append a 384x384 thumbnail of the whole image when more than one tile was produced. Then shrink the budget: starting from `max_num = 24`, while `sum_i (2 + tiles_i * 729) >= 3072`, set `max_num -= 2` when `max_num >= 3` else `max_num -= 1`, and retile. With one image this caps the total at 4 tiles (3 crops + thumbnail, 2918 tokens) or a single 384x384 tile.

3. Each `` in the templated text is replaced by `` + `` x (729 * tiles_i) + ``. The tokenizer prepends `<|plamo:bos|>` (1) (`add_bos_token: true`) and adds nothing at the end.

4. Each tile becomes `[3, 384, 384]` f32, `(pixel/255 - 0.5) / 0.5` per channel. `pixel_values: [total_tiles, 3, 384, 384]`.

### Decoder

The existing `Plamo2Model` with: RoPE base per layer = `rope_theta` if `layer_idx in full_attention_idx` else `rope_local_theta` (the two values are equal on both shipped checkpoints, so any layer gets 1e6); tied head when `tie_word_embeddings` is absent or true (ignore `lm_head.weight` in that case); EOS ids `[1]` from the top-level config.

## Implementation plan

1. **Detection.** `src/models/detection.rs`: add `"plamo2vl" => Ok(ModelType::Plamo2VLM)`. `src/models/mod.rs`: add `Plamo2VLM` to the `ModelType` enum, the supported-type list, the description table (`("PLaMo 2.1 VL (SigLIP tiles + MLP adapter + PLaMo 2 hybrid)", "Other VLM")`) and the `mlxcel list` grouping. Add a test in `src/models/detection_tests.rs` that a minimal `{"model_type":"plamo2vl","text_config":{...},"vision_config":{...}}` resolves to `Plamo2VLM` and that `{"model_type":"plamo2"}` still resolves to `Plamo2`.

2. **PLaMo 2 text-model gaps (`src/models/plamo2.rs`).**
- Add `rope_theta: f32` (serde default 10000.0), `rope_local_theta: Option` (defaults to `rope_theta`) and `full_attention_idx: Option>` to `ModelArgs`; in `from_weights` set `rope_base` per layer as described above instead of the literal at line 1024. Text-only `pfnet/plamo-2-1b` has none of these keys and keeps base 10000.
- Add `fn embed_tokens(&self, input_ids) -> Option>` returning `self.embed_tokens.forward(input_ids)`, and `fn forward_with_embeddings(...)` / `fn forward_with_embeddings_and_sequence_id(...)` that run `forward_with_caches` from a supplied `[B, T, hidden]` tensor. Factor the body of `forward_with_caches` so the embedding lookup is skipped when embeddings are given; the causal mask, Mamba conv/SSM state and `norm` / head logic are unchanged.
- Honor `tie_word_embeddings` when it is true (or absent): do not load `lm_head.weight` even if present. The bf16 VL export ships a tied copy under `lm_head.weight`; quantized exports drop it.
- Unit tests in `src/models/plamo2_tests.rs`: `rope_base_defaults_to_10000_without_keys`, `rope_base_uses_local_theta_for_non_full_attention_layers`, `forward_with_embeddings_matches_forward_on_token_embeddings` (tiny random config; inject `embed_tokens(ids)` and assert byte-identical logits).

3. **Tokenizer (`src/tokenizer/mod.rs`, `build_plamo_tokenizer`).** After the four fixed specials, read `tokenizer_config.json` `added_tokens_decoder` and append every entry whose id is `>= vocab.len()` as a special added token (`plamo_added_token(id, content)`); fall back to `additional_special_tokens` assigned sequentially from `vocab.len()` when `added_tokens_decoder` is missing. Assert the ids are contiguous with the jsonl length so `` = 100002 holds. Test in the tokenizer tests: a temp dir with a 5-row jsonl plus a config declaring `` at id 7 encodes `"ab"` as `[.., 7, ..]` and decodes it back with `skip_special_tokens`.

4. **Config structs (`src/vision/plamo2_vl_config.rs`, new).** `Plamo2VlVisionConfig` deserialized from `vision_config` with `#[serde(rename = "vision_encoder_hidden_size")]` style aliases for the eight prefixed keys plus `image_token_id` (default 100002), `image_feature_size` (1152), `image_proj_hidden_size`. Provide `fn to_siglip_config(&self) -> vision::config::VisionConfig` (model_type `"siglip_vision_model"`, `hidden_act` from `vision_encoder_hidden_act`). `Plamo2VlProcessorConfig` from `preprocessor_config.json`: `size.width`, `patch_size`, `downsample_ratio`, `tile_context_length`, `image_mean`, `image_std`.

5. **Processor (`src/vision/processors/plamo2_vl.rs`, new).** `Plamo2VlProcessor { image_size: 384, patch_size: 14, downsample_ratio: 1.0, tile_context_length: 3072, max_tiles: 24, mean, std }` with `fn preprocess_with_tiles(&self, images) -> (UniquePtr /*[total,3,S,S]*/, Vec /*tiles per image*/)` implementing the candidate-grid rule, the ratio-factor selection (not the InternVL absolute-difference rule), the thumbnail rule, and the context-budget shrink loop. Reuse `image::imageops::FilterType::CatmullRom` for the resize and the CHW normalize loop from `InternVLProcessor::append_normalized_chw` (make that helper `pub(crate)` or copy it with `mean/std = 0.5`). If `downsample_ratio != 1.0`, round `image_size` up to a multiple of `patch_size * (1/downsample_ratio)` and zero-pad each tile to that stride; shipped checkpoints use 1.0 so this branch is config-driven only. Tests in `src/vision/processors/plamo2_vl_tests.rs`: `square_image_yields_one_tile_no_thumbnail`, `wide_image_2x1_adds_thumbnail` (a 1000x300 image gives `(2,1)` + thumbnail = 3 tiles), `context_budget_caps_total_tiles` (a 4:3 image that would pick `(2,2)` under `max_num = 24` must end at a budget where `2 + tiles*729 < 3072`, so at most 4 tiles), `pixel_values_shape_and_normalization` (a solid 128-gray image maps to `(128/255 - 0.5)/0.5` in every channel).

6. **Adapter (`src/vision/connectors/plamo2_vl.rs`, new).** `Plamo2VlImageProjector { norm0: RMSNorm (offset 1.0 applied at load by adding 1.0 to the weight, eps = text rms_norm_eps), bias0, linear1: UnifiedLinear, bias1, linear2: UnifiedLinear, bias2 }` with `from_weights(weights, "image_proj", group_size, bits)` reading `image_proj.norm0.weight`, `image_proj.bias{0,1,2}._bias`, `image_proj.linear{1,2}.weight` (+ `.scales/.biases` when quantized) and `forward(x) -> [N, hidden]` using `mlxcel_core::gelu` (exact erf). Test `projector_forward_shape_and_offset_norm`: with `norm0.weight = 0` the norm is a plain unit-weight RMSNorm; check against a hand computation on a 2x4 input.

7. **Runtime (`src/vision/plamo2_vl.rs`, new).** `Plamo2VlModel { text_model: Plamo2Model, vision_model: SigLipVisionModel, projector: Plamo2VlImageProjector, processor: Plamo2VlProcessor, image_token_id: 100002, img_start_token_id: 100000, img_end_token_id: 100001, num_image_token: 729, eos_token_ids: vec![1] }`. `get_input_embeddings(input_ids, pixel_values)`: transpose `[N,3,S,S]` to NHWC, cast to the tower dtype, `vision_model.forward` (no feature-layer selection; take the post-layernorm output), reshape to `[N*729, 1152]`, project, then `merge::merge_llava(image_token_id, &features, &inputs_embeds, input_ids)`. Implement `LanguageModel` by delegation to `text_model` exactly as `InternVLChatVLM` does (`src/vision/internvl.rs:205-320`), keeping `supports_batching() == false` (Mamba state) and the sequence-state methods.

8. **Loader (`src/loading/vlm_plamo2_vl.rs`, new; registered in `src/loading/vlm.rs` and dispatched from `src/loading/mod.rs` for `ModelType::Plamo2VLM`).** Read the sanitized config; parse `text_config` into `plamo2::ModelArgs` after injecting the top-level `eos_token_id` (1) and the top-level `quantization` block when the text config has none; build the tower and the projector first from `&weights`, then hand the map to `Plamo2Model::from_weights(args, weights)` (it takes ownership; PLaMo keys are top-level `model.*` and the extra `vision_model.*` / `image_proj.*` keys are ignored by its keyed lookups); the tower is `SigLipVisionModel::from_weights_with_quant_and_gelu(&weights, &siglip_cfg, "vision_model.vision_encoder.model", gs, bits, false)`; skip `vision_model.vision_encoder.model.head.*`. Sanitize rules: `patch_embedding.weight` is `[1152, 3, 14, 14]` in HF exports and `[1152, 14, 14, 3]` in converted ones; transpose `[0,2,3,1]` only when `shape[1] == 3` (`conv_channels_last` in `src/vision/encoders/deepseekocr_sam.rs:373` is the existing gate); `model.layers.layers.N.mixer.conv1d.weight` goes through the existing `plamo2::sanitize_weights`. bf16 to f16 conversion follows `load_vlm_weights_common` (keep `.scales`/`.biases` bf16 on quantized exports).

9. **Dispatch.** `src/loaded_model.rs`: `LoadedModel::Plamo2VLM(vision::Plamo2VlModel)` plus the delegation macro arm. `src/loaded_model_capabilities.rs`: `VlmRuntimeRef::Plamo2Vl(&Plamo2VlModel)` and the `vlm_runtime()` mapping. `src/multimodal/vlm_runtime.rs`: a `VlmRuntimeRef::Plamo2Vl` arm that (a) renders the fixed instruction template around the user text with `images.len()` placeholders unless the prompt already contains `### 指示:`, (b) tokenizes with BOS, (c) calls `processor.preprocess_with_tiles`, (d) expands every 100002 placeholder in the token stream into `[100000] + [100002] * (729 * tiles_i) + [100001]` (generalize `insert_internvl_image_tokens` in `src/multimodal/internvl_prompt.rs` to take the three ids, it already does), (e) calls `get_input_embeddings`, returning `VlmPreparationSummary::Plamo2Vl { image_blocks, total_image_tokens }`; print it in `src/commands/generate_vlm.rs` next to the InternVL arm (line 228). Server: `/v1/chat/completions` image parts reach `prepare_vlm_inputs` through `src/server/chat_request.rs` unchanged; the PLaMo prompt builder must take the flattened user text (the checkpoint ships no Jinja chat template, so the fixed template replaces `apply_chat_template`).

10. **Output suppression.** Override `output_suppressed_token_ids` to `[100000, 100001, 100002]` so the placeholder ids can never be sampled (see the rationale at `src/lib/mlxcel-core/src/generate.rs:349-371`).

11. **Docs.** `docs/supported-models.md`: add PLaMo 2.1 VL under VLMs and drop the stale sentence in the PLaMo 2 entry (line 58) claiming the tokenizer loader does not read `tokenizer.jsonl`.

## Validation

(a) Unit tests listed above in `plamo2_tests.rs`, `plamo2_vl_tests.rs` (processor), `connectors/plamo2_vl_tests.rs`, tokenizer tests, `detection_tests.rs`.

(b) Real checkpoint: `pfnet/plamo-2.1-2b-vl` (bf16, 11.5 GB safetensors, 3.09B params; listed on the model card). Community 4-bit MLX conversions exist on the Hub (`tokimoa/plamo-2.1-2b-vl-mlx-4bit`, about 2 GB) and are the quick path; if they fail to load, quantize the bf16 original locally.

```
magick -size 512x512 xc:red /tmp/solid_red.png # any solid-color PNG works; tests/fixtures/test_image.png is the shared 224x224 photo
./target/release/mlxcel generate -m models/plamo-2.1-2b-vl --image /tmp/solid_red.png -p "この画像の色は何色ですか。" -n 32
./target/release/mlxcel generate -m models/plamo-2.1-2b-vl --image tests/fixtures/test_image.png -p "Describe this image." -n 64
```

Acceptance: finite logits at every step; the answer names red (赤 / red); the printed preparation summary shows 1 image block of 729 tokens for a square image and `3 * 729` or `4 * 729` for a 16:9 photo; a prompt with a 1000x300 image reports 3 tiles. Token-exact comparison against the checkpoint's own `transformers` implementation (`modeling_plamo2_vl.py`, CPU bf16, greedy) on the 2B model for a 64-token continuation of a real photo prompt; allow a divergence only after the first 32 tokens and only if both continuations stay fluent. Also run `pfnet/plamo-2-1b` text generation before and after step 2 to confirm the text-only path is byte-identical.

## Acceptance criteria

- [ ] `model_type: "plamo2vl"` resolves to `ModelType::Plamo2VLM`; `plamo2` still resolves to `Plamo2`
- [ ] `Plamo2Model` exposes `embed_tokens` and `forward_with_embeddings`, with a byte-identity test against the token path
- [ ] RoPE base is read from `rope_theta` / `rope_local_theta` / `full_attention_idx`; `pfnet/plamo-2-1b` output is unchanged
- [ ] `build_plamo_tokenizer` emits id 100002 for `` on the VL checkpoint
- [ ] Tiling, budget shrink, and `` framing match the rules above (unit tests)
- [ ] `mlxcel generate --image` and `/v1/chat/completions` with `image_url` produce a correct color answer on a solid image with the 2B checkpoint
- [ ] docs/supported-models.md updated
- [ ] detection table in src/models/detection.rs updated with a test
- [ ] cargo test --workspace --profile test-fast --features metal,accelerate passes
- [ ] cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass

## Out of scope

- Multi-image prompts beyond what the tile budget allows (the model is single-image by design; multiple `--image` values are accepted but share the 3072-token tile budget).
- Sliding-window attention (`attention_window_size`); the decoder keeps a full `KVCache`.
- Batched serving (`supports_batching` stays false, as for text PLaMo 2).

Contributor guide

Open the contributing guide

Research direction

Start by reading the existing PLaMo 2 path in src/models/plamo2.rs, detection in src/models/detection.rs, and the SigLIP and InternVL implementations named in the issue. Then follow the implementation plan through tokenizer/mod.rs and the new vision config and processor files, using the specified unit tests as checkpoints. Done means detection, tokenization, tiling, embedding injection, and both generate and chat image flows work for the listed checkpoints.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
computer-vision, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.