diff --git a/examples/qwen3_vlm/README.md b/examples/qwen3_vlm/README.md new file mode 100644 index 0000000..b508782 --- /dev/null +++ b/examples/qwen3_vlm/README.md @@ -0,0 +1,123 @@ +# Vision-Language Model with a Qwen3 backbone + +A worked example: build KempnerForge VLMs (Joint-Decoder, Cross-Attention, MoT, +MoMa) on a **Qwen3-0.6B** backbone, optionally **warm-start** them from an +externally-trained (`KempnerInstitute/multimodal`) checkpoint, and **fine-tune +on video** (WebVid-10M). + +The Qwen3 backbone is expressed **entirely through config** — there is no +Qwen-specific model class. It relies on a few small, *general* additions to +KempnerForge core (reusable by any model, not just this example): + +- `ModelConfig.head_dim` — an overridable field, so the attention width can + decouple from the model dim (Qwen3-0.6B: `dim=1024`, `head_dim=128` → + 2048-wide attention). Default preserves Llama-style `dim // n_heads`. +- `mlp_2layer` gains an optional `pre_norm` — a pre-projection norm (`ln_q`) + over the vision features, selected by norm-registry key (`rmsnorm` / + `layernorm`); off by default, so existing `mlp_2layer` configs are unchanged. +- Cross-attention `head_dim` threading — CA blocks honor the decoupled head_dim. +- `checkpoint.exclude_from_loading` is wired into the warm-start path in + `scripts/train.py`, so a weights-only checkpoint loads without an optimizer. + +Those live in `kempnerforge/`, tested alongside the components they extend +(`tests/unit/test_config.py`, `test_model.py`, `test_adapter.py`, +`test_cross_attention.py`). Everything specific to *this* use case lives here. + +## Contents + +``` +examples/qwen3_vlm/ +├── README.md +├── convert_multimodal_checkpoint.py # multimodal .pt -> KempnerForge DCP +├── configs/ +│ ├── vlm_qwen3_0.6b_joint_decoder.toml # warm-start from converted JD ckpt +│ ├── vlm_qwen3_0.6b_mot.toml # warm-start from converted MoT ckpt +│ ├── vlm_qwen3_0.6b_cross_attn.toml # from scratch (see note) +│ └── vlm_qwen3_0.6b_moma.toml # from scratch (no source ckpt) +└── tests/ + ├── test_configs.py # configs load + build the right Qwen3-0.6B shapes + └── test_converter.py # the key mapping (pure, no GPU/network) +``` + +## The four arches + +| arch | source checkpoint | this example | +|------|-------------------|--------------| +| `joint_decoder` | converts 1:1 | warm-start from a converted JD DCP | +| `mot` | converts 1:1, or init from a dense JD (duplicate per modality) | warm-start from a converted MoT DCP | +| `moma` | none (no multimodal-repo counterpart) | **from scratch** | +| `cross_attention` | exists, but doesn't convert faithfully yet | **from scratch** (see below) | + +**Cross-attention note.** The `multimodal` repo's cross-attention is an extra +sub-module at every 4th layer that projects fused K/V straight from the **vision +dim (1152)** and reuses that layer's FFN, with a vision-side norm and QK-norm. +KF's `CrossAttentionBlock` is a separate block that projects K/V from the +**text/adapter dim (1024)** with its own FFN and no vision-side norm. So the +trained cross-attn tensors don't map 1:1. A faithful CA checkpoint needs a small +KF cross-attn alignment (identity adapter + K/V from the vision dim + QK/vision +norms) — tracked as its own PR. Until then, cross-attention trains from scratch. + +## 1. Convert a checkpoint (optional; JD / MoT only) + +Network-free: builds the KF transformer + adapter on `meta` to validate +keys/shapes and carries the SigLIP tower straight from the source (no download). +Reads only the model `*.pt` (optimizer state is ignored) and refuses to write +inside the source tree. + +```bash +# Joint-Decoder (1:1) +uv run python examples/qwen3_vlm/convert_multimodal_checkpoint.py \ + --src /path/to/stage-1/vlm///model_epoch_0_100000.pt \ + --config examples/qwen3_vlm/configs/vlm_qwen3_0.6b_joint_decoder.toml \ + --out /your/scratch/converted/jd + +# Init a MoT checkpoint from a dense Joint-Decoder checkpoint (duplicate) +uv run python examples/qwen3_vlm/convert_multimodal_checkpoint.py \ + --src /path/to/jd_model.pt --from joint_decoder \ + --config examples/qwen3_vlm/configs/vlm_qwen3_0.6b_mot.toml \ + --out /your/scratch/converted/mot_from_jd + +# Native MoT source (1:1) — state which source modality index is the image stream +uv run python examples/qwen3_vlm/convert_multimodal_checkpoint.py \ + --src /path/to/mot_model.pt \ + --config examples/qwen3_vlm/configs/vlm_qwen3_0.6b_mot.toml \ + --out /your/scratch/converted/mot --mot-image-index 0 +``` + +Then point the config at the output: + +```toml +[checkpoint] +load_path = "/your/scratch/converted/jd" +exclude_from_loading = ["optimizer", "dataloader"] # weights-only warm start +``` + +## 2. Fine-tune on video (WebVid-10M) + +Each config wires the `[video]` pipeline (SigLIP2 so400m + a pre-norm +`mlp_2layer`, **4 frames × 256 tokens/frame** = 1024 visual tokens, vision +tower frozen). Video decoding needs the `video` extra (PyAV) — run +`uv sync --group video` before the first finetune. + +```bash +uv run torchrun --nproc_per_node=4 scripts/train.py \ + examples/qwen3_vlm/configs/vlm_qwen3_0.6b_joint_decoder.toml +``` + +Before a real run, edit the placeholders in the config: `[checkpoint].load_path` +(your converted DCP), `[video].data_root` (your WebVid corpus), `[checkpoint].dir`, +and the `[metrics]` wandb fields. SigLIP2 (`google/siglip2-so400m-patch14-224`) and +the `Qwen/Qwen3-0.6B` tokenizer must be reachable (local path or `HF_HOME`). + +**Frame budget.** The `mlp_2layer` connector keeps 256 tokens/frame (no pooling), +which is why these configs use 4 frames. For longer clips, switch `[adapter].type` +to a pooling connector (`avgpool` / `attentional_pool`) — that reduces tokens/frame +but re-initializes the adapter (the backbone + vision still warm-start). + +## 3. Run the tests + +Standalone (not part of the main `tests/` suite): + +```bash +uv run pytest examples/qwen3_vlm/tests/ +``` diff --git a/examples/qwen3_vlm/configs/vlm_qwen3_0.6b_cross_attn.toml b/examples/qwen3_vlm/configs/vlm_qwen3_0.6b_cross_attn.toml new file mode 100644 index 0000000..213ed69 --- /dev/null +++ b/examples/qwen3_vlm/configs/vlm_qwen3_0.6b_cross_attn.toml @@ -0,0 +1,93 @@ +# Video finetune: Qwen3-0.6B Cross-Attention VLM on WebVid-10M. +# +# Trains from scratch for now: a faithful warm-start needs a converted CA +# checkpoint, which requires a KF cross-attention arch-alignment first (the +# multimodal repo fuses K/V from the vision dim + shares the decoder FFN, vs +# KF's separate K/V from the text dim + a per-block FFN). Once that lands, set +# [checkpoint].load_path to the converted CA DCP. +# +# Image features flow as K/V into CrossAttentionBlocks (cadence 4 -> 7 blocks); +# the residual stream is text-only, so max_seq_len only needs to cover the text. +[model] +dim = 1024 +n_layers = 28 +n_heads = 16 +n_kv_heads = 8 +head_dim = 128 +vocab_size = 151936 +ffn_hidden_dim = 3072 +norm_type = "rmsnorm" +norm_eps = 1e-6 +activation = "silu" +qk_norm = true +rope_theta = 1000000.0 +tie_embeddings = false +max_seq_len = 512 + +[vision_encoder] +type = "siglip2" +path = "google/siglip2-so400m-patch14-224" + +[adapter] +type = "mlp_2layer" +pre_norm = "rmsnorm" +hidden_dim = 1152 +activation = "gelu" + +[vlm] +arch = "cross_attention" +max_text_len = 64 +cross_attention_every_n_layers = 4 +freeze = [{ module = "vision_encoder", frozen = true }] + +[video] +data_root = "/n/holylfs06/LABS/kempner_shared/Everyone/testbed/video/webvid-10m" +dataset_type = "webvid" +split = "train" +max_frames = 4 +min_frames = 2 +fps = 2.0 +frame_size = 224 +max_samples = 4096 + +[data] +tokenizer_path = "Qwen/Qwen3-0.6B" +num_workers = 4 +pin_memory = true + +[train] +batch_size = 4 +seq_len = 512 +max_steps = 5000 +grad_accum_steps = 1 +grad_clip_norm = 1.0 +seed = 42 +compile_model = false +activation_checkpointing = "none" + +[optimizer] +name = "adamw" +lr = 2e-5 +weight_decay = 0.1 +betas = [0.9, 0.95] +eps = 1e-8 +fused = true + +[scheduler] +name = "cosine" +warmup_steps = 100 +min_lr_ratio = 0.1 + +[checkpoint] +dir = "/n/netscratch/kempner_dev/amazloumi/video-arch/checkpoints/smoke/cross_video" +# load_path = "" # pending the CA arch-alignment +exclude_from_loading = ["optimizer", "dataloader"] +interval = 1000 +keep_last_n = 1 + +[metrics] +log_interval = 10 +enable_wandb = true +wandb_project = "kf-vlm-qwen3-video" +wandb_run_name = "qwen3-0.6b-cross-video" +enable_tensorboard = false diff --git a/examples/qwen3_vlm/configs/vlm_qwen3_0.6b_joint_decoder.toml b/examples/qwen3_vlm/configs/vlm_qwen3_0.6b_joint_decoder.toml new file mode 100644 index 0000000..a0aa348 --- /dev/null +++ b/examples/qwen3_vlm/configs/vlm_qwen3_0.6b_joint_decoder.toml @@ -0,0 +1,88 @@ +# Video finetune: Qwen3-0.6B Joint-Decoder, warm-started from a converted image +# checkpoint, on WebVid-10M. so400m + a pre-norm mlp_2layer match the converted +# checkpoint (256 tokens/frame, no pooling), so 4 frames * 256 = 1024 visual + +# 64 text fits max_seq_len. Set [checkpoint].load_path to your converted DCP and +# [video].data_root to your WebVid corpus. More frames need a pooling connector +# (avgpool/attentional_pool), which re-inits the adapter. +[model] +dim = 1024 +n_layers = 28 +n_heads = 16 +n_kv_heads = 8 +head_dim = 128 +vocab_size = 151936 +ffn_hidden_dim = 3072 +norm_type = "rmsnorm" +norm_eps = 1e-6 +activation = "silu" +qk_norm = true +rope_theta = 1000000.0 +tie_embeddings = false +max_seq_len = 1280 + +[vision_encoder] +type = "siglip2" +path = "google/siglip2-so400m-patch14-224" + +[adapter] +type = "mlp_2layer" +pre_norm = "rmsnorm" +hidden_dim = 1152 +activation = "gelu" + +[vlm] +arch = "joint_decoder" +max_text_len = 64 +freeze = [{ module = "vision_encoder", frozen = true }] + +[video] +data_root = "/n/holylfs06/LABS/kempner_shared/Everyone/testbed/video/webvid-10m" +dataset_type = "webvid" +split = "train" +max_frames = 4 +min_frames = 2 +fps = 2.0 +frame_size = 224 +max_samples = 4096 + +[data] +tokenizer_path = "Qwen/Qwen3-0.6B" +num_workers = 4 +pin_memory = true + +[train] +batch_size = 4 +seq_len = 1280 +max_steps = 5000 +grad_accum_steps = 1 +grad_clip_norm = 1.0 +seed = 42 +compile_model = false +activation_checkpointing = "none" + +[optimizer] +name = "adamw" +lr = 2e-5 +weight_decay = 0.1 +betas = [0.9, 0.95] +eps = 1e-8 +fused = true + +[scheduler] +name = "cosine" +warmup_steps = 100 +min_lr_ratio = 0.1 + +[checkpoint] +dir = "/n/netscratch/kempner_dev/amazloumi/video-arch/checkpoints/smoke/jd_video" +load_path = "/n/netscratch/kempner_dev/amazloumi/video-arch/checkpoints/coverted-old-stage1/joint-decoder/v06-si-f001-pre-s1" +exclude_from_loading = ["optimizer", "dataloader"] +interval = 1000 +keep_last_n = 1 + +[metrics] +log_interval = 10 +enable_wandb = true +wandb_project = "kf-vlm-qwen3-video" +wandb_run_name = "qwen3-0.6b-jd-video" +enable_tensorboard = false diff --git a/examples/qwen3_vlm/configs/vlm_qwen3_0.6b_moma.toml b/examples/qwen3_vlm/configs/vlm_qwen3_0.6b_moma.toml new file mode 100644 index 0000000..6bff196 --- /dev/null +++ b/examples/qwen3_vlm/configs/vlm_qwen3_0.6b_moma.toml @@ -0,0 +1,95 @@ +# Video finetune: Qwen3-0.6B MoMa VLM on WebVid-10M. +# +# Trains from scratch: MoMa has no multimodal-repo counterpart, so there is no +# stage-1 checkpoint to warm-start from. Shared Q/K/V/O attention + per-modality +# MoE FFN groups; image tokens prepend the text sequence (same residual layout +# as JD/MoT). MoMa v1 is training-only (expert-choice routing is non-causal). +[model] +dim = 1024 +n_layers = 28 +n_heads = 16 +n_kv_heads = 8 +head_dim = 128 +vocab_size = 151936 +ffn_hidden_dim = 3072 +norm_type = "rmsnorm" +norm_eps = 1e-6 +activation = "silu" +qk_norm = true +rope_theta = 1000000.0 +tie_embeddings = false +max_seq_len = 1280 + +[vision_encoder] +type = "siglip2" +path = "google/siglip2-so400m-patch14-224" + +[adapter] +type = "mlp_2layer" +pre_norm = "rmsnorm" +hidden_dim = 1152 +activation = "gelu" + +[vlm] +arch = "moma" +max_text_len = 64 +moma_gumbel_noise = true +freeze = [{ module = "vision_encoder", frozen = true }] + +[vlm.moma_experts_per_modality] +image = 2 +text = 2 + +[video] +data_root = "/n/holylfs06/LABS/kempner_shared/Everyone/testbed/video/webvid-10m" +dataset_type = "webvid" +split = "train" +max_frames = 4 +min_frames = 2 +fps = 2.0 +frame_size = 224 +max_samples = 4096 + +[data] +tokenizer_path = "Qwen/Qwen3-0.6B" +num_workers = 4 +pin_memory = true + +[train] +batch_size = 4 +seq_len = 1280 +max_steps = 5000 +grad_accum_steps = 1 +grad_clip_norm = 1.0 +seed = 42 +compile_model = false +activation_checkpointing = "none" + +[optimizer] +name = "adamw" +lr = 2e-5 +weight_decay = 0.1 +betas = [0.9, 0.95] +eps = 1e-8 +fused = true + +[scheduler] +name = "cosine" +warmup_steps = 100 +min_lr_ratio = 0.1 + +[distributed] +dp_shard = -1 +nccl_timeout_sec = 600 + +[checkpoint] +dir = "/n/netscratch/kempner_dev/amazloumi/video-arch/checkpoints/smoke/moma_video" +interval = 1000 +keep_last_n = 1 + +[metrics] +log_interval = 10 +enable_wandb = true +wandb_project = "kf-vlm-qwen3-video" +wandb_run_name = "qwen3-0.6b-moma-video" +enable_tensorboard = false diff --git a/examples/qwen3_vlm/configs/vlm_qwen3_0.6b_mot.toml b/examples/qwen3_vlm/configs/vlm_qwen3_0.6b_mot.toml new file mode 100644 index 0000000..cfe52ff --- /dev/null +++ b/examples/qwen3_vlm/configs/vlm_qwen3_0.6b_mot.toml @@ -0,0 +1,88 @@ +# Video finetune: Qwen3-0.6B MoT, warm-started from a converted image checkpoint +# (itself seeded from a dense Joint-Decoder), on WebVid-10M. so400m + +# a pre-norm mlp_2layer matches the converted checkpoint (256 tokens/frame), so +# 4 frames * 256 = 1024 visual + 64 text fits max_seq_len. Set +# [checkpoint].load_path to your converted MoT DCP and [video].data_root to your +# WebVid corpus. More frames need a pooling connector (re-inits the adapter). +[model] +dim = 1024 +n_layers = 28 +n_heads = 16 +n_kv_heads = 8 +head_dim = 128 +vocab_size = 151936 +ffn_hidden_dim = 3072 +norm_type = "rmsnorm" +norm_eps = 1e-6 +activation = "silu" +qk_norm = true +rope_theta = 1000000.0 +tie_embeddings = false +max_seq_len = 1280 + +[vision_encoder] +type = "siglip2" +path = "google/siglip2-so400m-patch14-224" + +[adapter] +type = "mlp_2layer" +pre_norm = "rmsnorm" +hidden_dim = 1152 +activation = "gelu" + +[vlm] +arch = "mot" +max_text_len = 64 +freeze = [{ module = "vision_encoder", frozen = true }] + +[video] +data_root = "/n/holylfs06/LABS/kempner_shared/Everyone/testbed/video/webvid-10m" +dataset_type = "webvid" +split = "train" +max_frames = 4 +min_frames = 2 +fps = 2.0 +frame_size = 224 +max_samples = 4096 + +[data] +tokenizer_path = "Qwen/Qwen3-0.6B" +num_workers = 4 +pin_memory = true + +[train] +batch_size = 4 +seq_len = 1280 +max_steps = 5000 +grad_accum_steps = 1 +grad_clip_norm = 1.0 +seed = 42 +compile_model = false +activation_checkpointing = "none" + +[optimizer] +name = "adamw" +lr = 2e-5 +weight_decay = 0.1 +betas = [0.9, 0.95] +eps = 1e-8 +fused = true + +[scheduler] +name = "cosine" +warmup_steps = 100 +min_lr_ratio = 0.1 + +[checkpoint] +dir = "/n/netscratch/kempner_dev/amazloumi/video-arch/checkpoints/smoke/mot_video" +load_path = "/n/netscratch/kempner_dev/amazloumi/video-arch/checkpoints/coverted-old-stage1/mot/v06-si-f001-pre-s1" +exclude_from_loading = ["optimizer", "dataloader"] +interval = 1000 +keep_last_n = 1 + +[metrics] +log_interval = 10 +enable_wandb = true +wandb_project = "kf-vlm-qwen3-video" +wandb_run_name = "qwen3-0.6b-mot-video" +enable_tensorboard = false diff --git a/examples/qwen3_vlm/convert_multimodal_checkpoint.py b/examples/qwen3_vlm/convert_multimodal_checkpoint.py new file mode 100644 index 0000000..6dd53e2 --- /dev/null +++ b/examples/qwen3_vlm/convert_multimodal_checkpoint.py @@ -0,0 +1,482 @@ +#!/usr/bin/env python3 +"""Convert a KempnerInstitute/multimodal Qwen3 VLM checkpoint to KempnerForge DCP. + +One-time import tool for **fine-tuning on top of** an externally-trained +(multimodal-repo) Qwen3-backbone VLM. It maps the source's flat ``torch.save`` +state dict onto KempnerForge's VLM state-dict layout — the transformer + adapter +are built from a target TOML (one of the ``configs/train/vlm_qwen3_0.6b_*.toml``) +to validate keys/shapes, and the SigLIP vision tower is carried straight from the +source (no vision download) — then writes a DCP directory the warm-start loads: + + [checkpoint] + load_path = "" + exclude_from_loading = ["optimizer", "dataloader"] # weights-only + +The mapping is derived purely from the checkpoint's tensor names + shapes and +KempnerForge's own target layout (a clean-room, data-only mapping — no import of +the source repo's code). ``map_key`` is a pure function so it is unit-tested on +key strings alone. + +Supported conversions ``(source arch -> target arch)``: + - ``joint_decoder -> joint_decoder`` — 1:1. + - ``mot -> mot`` — 1:1 per modality (needs ``--mot-image-index``). + - ``joint_decoder -> mot`` — **init MoT from a dense JD checkpoint** by + duplicating each per-layer dense weight into both modality copies (the + canonical MoT warm-start; no ``--mot-image-index`` since both copies are + identical). This mirrors how the multimodal repo built its MoT checkpoints + (they were never trained from scratch — always seeded from JD). + +``cross_attention`` differs structurally between the two repos (KF uses separate +K/V from the text dim + a per-block FFN; the source fuses K/V from the vision dim +and shares the decoder-layer FFN) and needs a KF CA arch-alignment first; +``moma`` has no source counterpart. Both raise a clear error. + +Optimizer state is never read — the source ``model_*.pt`` is a bare model state +dict, and its ``optimizer_*.pt`` sibling is ignored (fine-tuning starts fresh). + +Usage: + # Joint-Decoder (1:1) + uv run python scripts/convert_multimodal_checkpoint.py \ + --src /path/to/model_epoch_0_100000.pt \ + --config configs/train/vlm_qwen3_0.6b_joint_decoder.toml \ + --out /path/to/output/jd_dcp + + # Init a MoT checkpoint from a dense Joint-Decoder checkpoint + uv run python scripts/convert_multimodal_checkpoint.py \ + --src /path/to/jd_model.pt --from joint_decoder \ + --config configs/train/vlm_qwen3_0.6b_mot.toml \ + --out /path/to/output/mot_from_jd_dcp + + # Native MoT source (1:1) — state which source modality index is the image + uv run python scripts/convert_multimodal_checkpoint.py \ + --src /path/to/mot_model.pt \ + --config configs/train/vlm_qwen3_0.6b_mot.toml \ + --out /path/to/output/mot_dcp \ + --mot-image-index 0 +""" + +from __future__ import annotations + +import argparse +import logging +from pathlib import Path +from typing import Any + +import torch +import torch.distributed.checkpoint as dcp + +from kempnerforge.config.loader import load_config +from kempnerforge.model.adapter import build_adapter +from kempnerforge.model.transformer import Transformer + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +logger = logging.getLogger(__name__) + +# Supported (source arch, target arch) conversions. +SUPPORTED_CONVERSIONS: tuple[tuple[str, str], ...] = ( + ("joint_decoder", "joint_decoder"), + ("mot", "mot"), + ("joint_decoder", "mot"), # init MoT from a dense JD checkpoint (duplicate) +) + +# KF MoT modality names (KF MoTConfig default). JD->MoT duplicates each dense +# weight into both; the fanned-out final norm seeds both mot_norms. +MOT_MODALITIES = ("image", "text") + +# Source (multimodal repo) top-level prefixes. +_VISION_SRC_PREFIX = "image_encoder.model.vision_model." +_VISION_KF_PREFIX = "vision_encoder.vision_tower." +_CORE_SRC_PREFIX = "multimodal_core." + +# Adapter (qwen2_5_vl_patch_merger) suffix map: source ``adapter.model.`` +# -> KF ``adapter.``. ln_q is the merger's pre-projection RMSNorm; mlp.0 +# / mlp.2 are the two Linears (mlp.1 is the activation, no params). +_ADAPTER_SUFFIX = { + "ln_q.weight": "ln_q.weight", + "mlp.0.weight": "proj1.weight", + "mlp.0.bias": "proj1.bias", + "mlp.2.weight": "proj2.weight", + "mlp.2.bias": "proj2.bias", +} + +# Per-modality attention sub-projections shared by the MoT layout. +_MOT_ATTN_PROJ = ("q_proj", "k_proj", "v_proj", "o_proj", "q_norm", "k_norm") + + +# --------------------------------------------------------------------------- +# Pure key mapping (unit-tested; no I/O, no model) +# --------------------------------------------------------------------------- + + +def _map_jd_layer(body: str) -> str: + """Map a Joint-Decoder ``layers.{i}.`` body to KF naming (1:1).""" + body = body.replace(".self_attn.", ".attention.") + body = body.replace(".input_layernorm.", ".attention_norm.") + body = body.replace(".post_attention_layernorm.", ".mlp_norm.") + return body + + +def _map_mot_layer(body: str, idx_to_name: dict[str, str]) -> str: + """Map a native-MoT ``layers.{i}.`` body (per-modality index) to KF. + + Source ``{m}`` is a numeric modality index (0/1); KF keys it by name + (image/text). ``idx_to_name`` supplies that resolution. + """ + parts = body.split(".") + i = parts[1] + grp = parts[2] + if grp == "self_attn": + sub = parts[3] + if sub == "input_layer_norm": # self_attn.input_layer_norm.{m}.weight + name = idx_to_name[parts[4]] + return f"layers.{i}.attn_norm.{name}.{'.'.join(parts[5:])}" + if sub in _MOT_ATTN_PROJ: # self_attn.{proj}.{m}.weight + name = idx_to_name[parts[4]] + return f"layers.{i}.attn.{sub}.{name}.{'.'.join(parts[5:])}" + elif grp == "feed_forward": + sub = parts[3] + if sub == "mlp": # feed_forward.mlp.{m}.{proj}.weight + name = idx_to_name[parts[4]] + return f"layers.{i}.mlp.{name}.{'.'.join(parts[5:])}" + if sub == "post_attention_layernorm": # feed_forward.post_attention_layernorm.{m}.weight + name = idx_to_name[parts[4]] + return f"layers.{i}.mlp_norm.{name}.{'.'.join(parts[5:])}" + raise KeyError(f"unrecognized MoT layer key: multimodal_core.{body}") + + +def _jd_to_mot_layer_one(body: str, name: str) -> str: + """Map a dense JD ``layers.{i}.`` body to ONE MoT modality's key. + + The dense weight is copied into modality ``name`` (the caller emits one key + per modality, so the same dense tensor seeds both the image and text copies). + """ + parts = body.split(".") + i = parts[1] + grp = parts[2] + if grp == "self_attn": + sub = parts[3] # q_proj/k_proj/v_proj/o_proj/q_norm/k_norm + return f"layers.{i}.attn.{sub}.{name}.{'.'.join(parts[4:])}" + if grp == "input_layernorm": + return f"layers.{i}.attn_norm.{name}.{'.'.join(parts[3:])}" + if grp == "post_attention_layernorm": + return f"layers.{i}.mlp_norm.{name}.{'.'.join(parts[3:])}" + if grp == "mlp": + proj = parts[3] # gate_proj/up_proj/down_proj + return f"layers.{i}.mlp.{name}.{proj}.{'.'.join(parts[4:])}" + raise KeyError(f"unrecognized JD layer key: multimodal_core.{body}") + + +def map_key( + src_key: str, + source_arch: str, + mot_idx_to_name: dict[str, str] | None = None, + target_arch: str | None = None, +) -> list[str] | None: + """Map one source key to zero or more KempnerForge ``VLMWrapper`` keys. + + ``target_arch`` defaults to ``source_arch`` (a 1:1 conversion). Set it + differently for a cross-arch init, e.g. ``source_arch="joint_decoder"`` + + ``target_arch="mot"`` duplicates each dense layer weight into both modality + copies. + + Returns: + - ``[kf_key, ...]`` — target key(s). One source tensor can seed several + KF keys (MoT final norm -> shared ``norm`` + both ``mot_norms``; + JD->MoT layer weight -> both modality copies). + - ``[]`` — intentionally dropped (e.g. ``text_head.bias``: KF's output + head is bias-free). + - ``None`` — unrecognized (the caller surfaces it as unmapped). + """ + target = target_arch or source_arch + + # Vision tower: identical HF SiglipVisionTransformer submodule on both + # sides, so a straight prefix swap (all arches). + if src_key.startswith(_VISION_SRC_PREFIX): + return [_VISION_KF_PREFIX + src_key[len(_VISION_SRC_PREFIX) :]] + + if src_key == "text_preprocessor.embed.weight": + return ["transformer.token_embedding.embedding.weight"] + if src_key == "text_head.weight": + return ["transformer.output_head.proj.weight"] + if src_key == "text_head.bias": + return [] # KF OutputHead has no bias (current Qwen3 checkpoints match) + + if src_key.startswith("adapter.model."): + suffix = src_key[len("adapter.model.") :] + mapped = _ADAPTER_SUFFIX.get(suffix) + return [f"adapter.{mapped}"] if mapped is not None else None + + if src_key == "multimodal_core.norm.weight": + if target == "mot": + # Source has one final norm; KF MoT applies per-modality final + # norms (mot_norms) and keeps the unused shared ``norm`` too. + return ["transformer.norm.weight"] + [ + f"transformer.mot_norms.{m}.weight" for m in MOT_MODALITIES + ] + return ["transformer.norm.weight"] + + if src_key.startswith(_CORE_SRC_PREFIX): + body = src_key[len(_CORE_SRC_PREFIX) :] # e.g. "layers.0.self_attn.q_proj.weight" + if not body.startswith("layers."): + return None + if source_arch == "joint_decoder" and target == "joint_decoder": + return [f"transformer.{_map_jd_layer(body)}"] + if source_arch == "mot" and target == "mot": + assert mot_idx_to_name is not None, "mot->mot requires mot_idx_to_name" + return [f"transformer.{_map_mot_layer(body, mot_idx_to_name)}"] + if source_arch == "joint_decoder" and target == "mot": + # Duplicate the dense weight into every modality copy. + return [f"transformer.{_jd_to_mot_layer_one(body, m)}" for m in MOT_MODALITIES] + raise NotImplementedError(f"conversion {source_arch!r} -> {target!r} not supported") + + return None + + +def map_state_dict( + src_sd: dict[str, torch.Tensor], + source_arch: str, + mot_idx_to_name: dict[str, str] | None = None, + target_arch: str | None = None, +) -> tuple[dict[str, torch.Tensor], list[str]]: + """Apply ``map_key`` across a source state dict. + + Returns ``(converted, unmapped)`` where ``converted`` is the KF-keyed state + dict (one source tensor may seed several KF keys) and ``unmapped`` lists any + source keys with no mapping (should be empty for a clean source). + """ + converted: dict[str, torch.Tensor] = {} + unmapped: list[str] = [] + for src_key, tensor in src_sd.items(): + targets = map_key(src_key, source_arch, mot_idx_to_name, target_arch) + if targets is None: + unmapped.append(src_key) + continue + for kf_key in targets: + converted[kf_key] = tensor + return converted, unmapped + + +# --------------------------------------------------------------------------- +# Shell: load source, build target, validate, write DCP +# --------------------------------------------------------------------------- + +_WRAPPER_KEYS = ("model", "model_state_dict", "state_dict") + + +def load_source_state_dict(path: Path) -> dict[str, torch.Tensor]: + """mmap-load a source ``model_*.pt`` and return its flat tensor state dict.""" + try: + obj = torch.load(path, map_location="cpu", mmap=True, weights_only=True) + except Exception as e: # noqa: BLE001 -- trusted internal ckpt; retry unsafe + logger.warning("weights_only load failed (%s); retrying weights_only=False", e) + obj = torch.load(path, map_location="cpu", mmap=True, weights_only=False) + if isinstance(obj, dict) and not any(torch.is_tensor(v) for v in obj.values()): + for k in _WRAPPER_KEYS: + if k in obj and isinstance(obj[k], dict): + logger.info("Unwrapped source state dict from ['%s']", k) + obj = obj[k] + break + if not isinstance(obj, dict): + raise TypeError(f"expected a state-dict mapping in {path}, got {type(obj)}") + return {k: v for k, v in obj.items() if torch.is_tensor(v)} + + +def _assert_out_not_in_source_tree(src: Path, out: Path) -> None: + """Refuse to write anywhere inside the source checkpoint's directory.""" + src_dir = src.resolve().parent + out_r = out.resolve() + if out_r == src_dir or src_dir in out_r.parents: + raise ValueError( + f"refusing to write output into the source checkpoint tree ({src_dir}). " + "Choose an --out path outside it; the reference checkpoints must stay intact." + ) + + +def _infer_feature_dim(src_sd: dict[str, torch.Tensor], config: Any) -> int: + """Vision feature dim used to size the adapter for the exact shape check. + + Prefer the source adapter's ``ln_q`` (the merger's pre-norm over the vision + feature dim); fall back to the config's ``vision_encoder.feature_dim`` when set. + """ + ln_q = src_sd.get("adapter.model.ln_q.weight") + if ln_q is not None: + return int(ln_q.shape[0]) + if config.vision_encoder is not None and config.vision_encoder.feature_dim > 0: + return int(config.vision_encoder.feature_dim) + raise ValueError( + "cannot infer vision feature_dim: source has no adapter.model.ln_q.weight and " + "vision_encoder.feature_dim is 0. Set feature_dim in the config." + ) + + +def convert( + src: str, + config_path: str, + out: str, + from_arch: str | None = None, + mot_image_index: int | None = None, + dtype: torch.dtype | None = None, + allow_partial: bool = False, +) -> None: + """Convert one multimodal ``.pt`` into a KempnerForge DCP directory.""" + src_path = Path(src) + out_path = Path(out) + _assert_out_not_in_source_tree(src_path, out_path) + + config = load_config(config_path, cli_args=[]) + target_arch = config.vlm.arch if config.vlm is not None else None + source_arch = from_arch or target_arch + if (source_arch, target_arch) not in SUPPORTED_CONVERSIONS: + raise NotImplementedError( + f"unsupported conversion {source_arch!r} -> {target_arch!r}. " + f"Supported: {sorted(SUPPORTED_CONVERSIONS)}. cross_attention needs a KF " + "CA arch-alignment; moma has no source counterpart." + ) + assert source_arch is not None and target_arch is not None # both in SUPPORTED_CONVERSIONS + + # Native MoT source carries numeric modality indices, so we must know which + # is the image stream. JD->MoT duplicates identical copies, so order is moot. + mot_idx_to_name: dict[str, str] | None = None + if source_arch == "mot" and target_arch == "mot": + if mot_image_index not in (0, 1): + raise ValueError( + "mot->mot requires --mot-image-index {0,1}: which source modality " + "index is the IMAGE stream. Verify against your training convention " + "(a wrong value silently swaps the image/text experts)." + ) + text_index = 1 - mot_image_index + mot_idx_to_name = {str(mot_image_index): "image", str(text_index): "text"} + logger.warning( + "MoT modality order: source index %d -> image, %d -> text. " + "Confirm this matches how the source was trained.", + mot_image_index, + text_index, + ) + + logger.info("Loading source checkpoint: %s", src_path) + src_sd = load_source_state_dict(src_path) + logger.info("Source tensors: %d (%s -> %s)", len(src_sd), source_arch, target_arch) + + assert config.vlm is not None # narrowed by the SUPPORTED_CONVERSIONS check above + if config.adapter is None: + raise ValueError("VLM config must define an [adapter] section") + + # Expected KF transformer + adapter keys/shapes from a meta build. No vision + # download: the SigLIP tower is carried straight from the source checkpoint + # (same HF submodule KF wraps), so building it here would only be discarded. + feature_dim = _infer_feature_dim(src_sd, config) + with torch.device("meta"): + transformer = Transformer(config.model, vlm_config=config.vlm, num_image_tokens=256) + adapter = build_adapter(config.adapter, in_dim=feature_dim, out_dim=config.model.dim) + expected: dict[str, tuple[int, ...]] = { + f"transformer.{k}": tuple(v.shape) for k, v in transformer.state_dict().items() + } + expected |= {f"adapter.{k}": tuple(v.shape) for k, v in adapter.state_dict().items()} + + converted, unmapped = map_state_dict(src_sd, source_arch, mot_idx_to_name, target_arch) + if dtype is not None: + converted = {k: v.to(dtype) for k, v in converted.items()} + + vision = {k for k in converted if k.startswith(_VISION_KF_PREFIX)} + nonvision = converted.keys() - vision + + missing = expected.keys() - nonvision + unexpected = nonvision - expected.keys() + if unmapped: + logger.warning("Unmapped source keys (%d): %s", len(unmapped), sorted(unmapped)[:10]) + if unexpected: + raise RuntimeError( + f"{len(unexpected)} converted keys are not in the target model (mapping bug): " + f"{sorted(unexpected)[:10]}" + ) + if missing: + msg = f"{len(missing)} target keys unfilled by the source: {sorted(missing)[:10]}" + if allow_partial: + logger.warning("%s (continuing: --allow-partial; these keep init)", msg) + else: + raise RuntimeError(msg + " (pass --allow-partial to fill them from init)") + if not vision: + logger.warning("No vision_encoder.* keys carried from the source (expected the tower).") + + # Shape check on the transformer/adapter keys (catches a wrong-shape slip). + bad = [ + (k, tuple(converted[k].shape), expected[k]) + for k in nonvision + if tuple(converted[k].shape) != expected[k] + ] + if bad: + raise RuntimeError(f"shape mismatches (source vs target): {bad[:8]}") + + out_path.mkdir(parents=True, exist_ok=True) + dcp.save({"model": converted}, checkpoint_id=str(out_path)) + logger.info( + "Wrote DCP (%d tensors: %d transformer+adapter, %d vision; %s -> %s) -> %s\n" + "Load it via [checkpoint].load_path + " + "exclude_from_loading=['optimizer','dataloader'] (weights-only warm start).", + len(converted), + len(nonvision), + len(vision), + source_arch, + target_arch, + out_path, + ) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0] if __doc__ else None) + parser.add_argument("--src", required=True, help="Source multimodal model_*.pt file") + parser.add_argument("--config", required=True, help="Target KF TOML (vlm_qwen3_0.6b_*.toml)") + parser.add_argument( + "--out", required=True, help="Output DCP directory (outside the source tree)" + ) + parser.add_argument( + "--from", + dest="from_arch", + default=None, + choices=["joint_decoder", "mot"], + help="Source arch (default: same as the target config's arch). " + "Use --from joint_decoder with a MoT config to init MoT from a dense JD checkpoint.", + ) + parser.add_argument( + "--mot-image-index", + type=int, + default=None, + choices=[0, 1], + help="For a native mot->mot conversion: which source modality index is the image stream", + ) + parser.add_argument( + "--dtype", + default=None, + choices=["float32", "bfloat16", "float16"], + help="Optional cast for the written weights (default: keep source dtype)", + ) + parser.add_argument( + "--allow-partial", + action="store_true", + help="Warn instead of failing when some target keys are unfilled (they keep init)", + ) + args = parser.parse_args(argv) + + dtype_map: dict[str, Any] = { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + "float16": torch.float16, + } + convert( + src=args.src, + config_path=args.config, + out=args.out, + from_arch=args.from_arch, + mot_image_index=args.mot_image_index, + dtype=dtype_map.get(args.dtype) if args.dtype else None, + allow_partial=args.allow_partial, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/qwen3_vlm/tests/test_configs.py b/examples/qwen3_vlm/tests/test_configs.py new file mode 100644 index 0000000..4bd979c --- /dev/null +++ b/examples/qwen3_vlm/tests/test_configs.py @@ -0,0 +1,108 @@ +"""Standalone tests for the Qwen3-0.6B VLM example configs. + +Not part of the main ``tests/`` suite (testpaths). Run explicitly: + + uv run pytest examples/qwen3_vlm/tests/ + +Each config loads, selects the right arch, builds a Qwen3-0.6B backbone block +with the decoupled 2048-wide attention, uses the SigLIP2 so400m encoder + +a pre-norm mlp_2layer adapter, and wires the video pipeline. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from kempnerforge.config.loader import load_config +from kempnerforge.model.transformer import TransformerBlock + +EXAMPLE = Path(__file__).resolve().parents[1] +ATTN_WIDTH = 16 * 128 # 2048, decoupled from dim=1024 + +QWEN3_06B = { + "dim": 1024, + "n_layers": 28, + "n_heads": 16, + "n_kv_heads": 8, + "head_dim": 128, + "ffn_hidden_dim": 3072, + "vocab_size": 151936, + "rope_theta": 1000000.0, + "norm_eps": 1e-6, +} + +CONFIGS = { + "joint_decoder": EXAMPLE / "configs/vlm_qwen3_0.6b_joint_decoder.toml", + "cross_attention": EXAMPLE / "configs/vlm_qwen3_0.6b_cross_attn.toml", + "mot": EXAMPLE / "configs/vlm_qwen3_0.6b_mot.toml", + "moma": EXAMPLE / "configs/vlm_qwen3_0.6b_moma.toml", +} + + +class TestQwen3VLMConfigs: + @pytest.mark.parametrize("arch", list(CONFIGS)) + def test_loads_and_selects_arch(self, arch): + cfg = load_config(str(CONFIGS[arch]), cli_args=[]) + assert cfg.vlm is not None + assert cfg.vlm.arch == arch + + @pytest.mark.parametrize("arch", list(CONFIGS)) + def test_model_is_qwen3_06b_shape(self, arch): + model = load_config(str(CONFIGS[arch]), cli_args=[]).model + assert model.dim == QWEN3_06B["dim"] + assert model.n_layers == QWEN3_06B["n_layers"] + assert model.n_heads == QWEN3_06B["n_heads"] + assert model.n_kv_heads == QWEN3_06B["n_kv_heads"] + assert model.head_dim == QWEN3_06B["head_dim"] + assert model.computed_ffn_hidden_dim == QWEN3_06B["ffn_hidden_dim"] + assert model.vocab_size == QWEN3_06B["vocab_size"] + assert model.qk_norm is True + assert model.tie_embeddings is False + assert model.rope_theta == QWEN3_06B["rope_theta"] + assert model.norm_eps == QWEN3_06B["norm_eps"] + + @pytest.mark.parametrize("arch", list(CONFIGS)) + def test_uses_prenorm_mlp_and_siglip2(self, arch): + cfg = load_config(str(CONFIGS[arch]), cli_args=[]) + assert cfg.adapter is not None + assert cfg.vision_encoder is not None + assert cfg.adapter.type == "mlp_2layer" + assert cfg.adapter.pre_norm == "rmsnorm" + assert cfg.vision_encoder.type == "siglip2" + assert "siglip2-so400m-patch14-224" in cfg.vision_encoder.path + # feature_dim / num_tokens use the 0 = "probe at build time" sentinel. + assert cfg.vision_encoder.feature_dim == 0 + assert cfg.vision_encoder.num_tokens == 0 + + @pytest.mark.parametrize("arch", list(CONFIGS)) + def test_video_pipeline_configured(self, arch): + cfg = load_config(str(CONFIGS[arch]), cli_args=[]) + assert cfg.video is not None + assert cfg.video.dataset_type == "webvid" + assert cfg.video.max_frames >= 1 + + @pytest.mark.parametrize("arch", ["joint_decoder", "mot"]) + def test_warm_start_configs_set_load_path(self, arch): + # JD and MoT warm-start from a converted checkpoint; cross/moma are + # from scratch (no source) and intentionally leave load_path unset. + cfg = load_config(str(CONFIGS[arch]), cli_args=[]) + assert cfg.checkpoint.load_path + assert cfg.checkpoint.exclude_from_loading == ["optimizer", "dataloader"] + + @pytest.mark.parametrize("arch", ["cross_attention", "moma"]) + def test_from_scratch_configs_have_no_load_path(self, arch): + cfg = load_config(str(CONFIGS[arch]), cli_args=[]) + assert not cfg.checkpoint.load_path + + @pytest.mark.parametrize("arch", list(CONFIGS)) + def test_backbone_block_builds_with_qwen3_attention_shape(self, arch): + # Build ONE block from the parsed config (cheap: no embedding/head, no + # full 28-layer stack) and confirm the decoupled 2048-wide attention. + model = load_config(str(CONFIGS[arch]), cli_args=[]).model + block = TransformerBlock(model, layer_idx=0) + assert block.attention.q_proj.out_features == ATTN_WIDTH # 2048 + assert block.attention.o_proj.out_features == QWEN3_06B["dim"] # 1024 + assert block.attention.q_norm is not None # qk_norm on + assert block.attention.q_norm.weight.shape == (QWEN3_06B["head_dim"],) # (128,) diff --git a/examples/qwen3_vlm/tests/test_converter.py b/examples/qwen3_vlm/tests/test_converter.py new file mode 100644 index 0000000..cf3bee6 --- /dev/null +++ b/examples/qwen3_vlm/tests/test_converter.py @@ -0,0 +1,362 @@ +"""Standalone tests for the multimodal -> KempnerForge checkpoint key mapping. + +Not part of the main ``tests/`` suite. Run explicitly: + + uv run pytest examples/qwen3_vlm/tests/ + +The mapping (``map_key`` / ``map_state_dict``) is pure, so it is tested on key +strings alone. The completeness test builds a tiny KF ``Transformer`` + adapter +(no vision, no network) for Joint-Decoder and MoT, synthesizes the matching +multimodal source keys, and asserts the mapping covers every target key exactly. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +import torch + +# Add the example dir to path so we can import the conversion module. +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from convert_multimodal_checkpoint import map_key, map_state_dict + +from kempnerforge.config.adapter import AdapterConfig +from kempnerforge.config.schema import ModelConfig +from kempnerforge.config.vlm import JointDecoderConfig, MoTConfig +from kempnerforge.model.adapter import build_adapter +from kempnerforge.model.transformer import Transformer + +IMG_TEXT = {"0": "image", "1": "text"} # source index -> KF modality name + + +# --------------------------------------------------------------------------- +# map_key — shared / vision / embed / head / adapter +# --------------------------------------------------------------------------- + + +class TestMapKeyShared: + def test_vision_prefix_swap(self): + assert map_key( + "image_encoder.model.vision_model.embeddings.patch_embedding.weight", "joint_decoder" + ) == ["vision_encoder.vision_tower.embeddings.patch_embedding.weight"] + assert map_key( + "image_encoder.model.vision_model.encoder.layers.7.self_attn.k_proj.bias", "mot" + ) == ["vision_encoder.vision_tower.encoder.layers.7.self_attn.k_proj.bias"] + + def test_token_embedding(self): + assert map_key("text_preprocessor.embed.weight", "joint_decoder") == [ + "transformer.token_embedding.embedding.weight" + ] + + def test_output_head_weight(self): + assert map_key("text_head.weight", "joint_decoder") == [ + "transformer.output_head.proj.weight" + ] + + def test_output_head_bias_dropped(self): + # KF OutputHead is bias-free; the source bias (older ckpts) is dropped. + assert map_key("text_head.bias", "joint_decoder") == [] + + def test_adapter_patch_merger(self): + assert map_key("adapter.model.ln_q.weight", "joint_decoder") == ["adapter.ln_q.weight"] + assert map_key("adapter.model.mlp.0.weight", "mot") == ["adapter.proj1.weight"] + assert map_key("adapter.model.mlp.0.bias", "mot") == ["adapter.proj1.bias"] + assert map_key("adapter.model.mlp.2.weight", "joint_decoder") == ["adapter.proj2.weight"] + assert map_key("adapter.model.mlp.2.bias", "joint_decoder") == ["adapter.proj2.bias"] + + def test_unrecognized_key_returns_none(self): + assert map_key("something.unexpected", "joint_decoder") is None + assert map_key("multimodal_core.mystery.weight", "joint_decoder") is None + + +# --------------------------------------------------------------------------- +# map_key — Joint-Decoder layers + final norm +# --------------------------------------------------------------------------- + + +class TestMapKeyJointDecoder: + def test_attention_projections(self): + p = "multimodal_core.layers.5.self_attn" + assert map_key(f"{p}.q_proj.weight", "joint_decoder") == [ + "transformer.layers.5.attention.q_proj.weight" + ] + assert map_key(f"{p}.o_proj.weight", "joint_decoder") == [ + "transformer.layers.5.attention.o_proj.weight" + ] + assert map_key(f"{p}.q_norm.weight", "joint_decoder") == [ + "transformer.layers.5.attention.q_norm.weight" + ] + + def test_layernorms(self): + assert map_key("multimodal_core.layers.0.input_layernorm.weight", "joint_decoder") == [ + "transformer.layers.0.attention_norm.weight" + ] + assert map_key( + "multimodal_core.layers.0.post_attention_layernorm.weight", "joint_decoder" + ) == ["transformer.layers.0.mlp_norm.weight"] + + def test_mlp(self): + assert map_key("multimodal_core.layers.2.mlp.gate_proj.weight", "joint_decoder") == [ + "transformer.layers.2.mlp.gate_proj.weight" + ] + + def test_final_norm_single(self): + assert map_key("multimodal_core.norm.weight", "joint_decoder") == [ + "transformer.norm.weight" + ] + + +# --------------------------------------------------------------------------- +# map_key — MoT per-modality + fanned-out final norm +# --------------------------------------------------------------------------- + + +class TestMapKeyMoT: + def test_attn_norm_per_modality(self): + assert map_key( + "multimodal_core.layers.1.self_attn.input_layer_norm.0.weight", "mot", IMG_TEXT + ) == ["transformer.layers.1.attn_norm.image.weight"] + assert map_key( + "multimodal_core.layers.1.self_attn.input_layer_norm.1.weight", "mot", IMG_TEXT + ) == ["transformer.layers.1.attn_norm.text.weight"] + + def test_attn_projections_per_modality(self): + assert map_key("multimodal_core.layers.3.self_attn.q_proj.0.weight", "mot", IMG_TEXT) == [ + "transformer.layers.3.attn.q_proj.image.weight" + ] + assert map_key("multimodal_core.layers.3.self_attn.k_norm.1.weight", "mot", IMG_TEXT) == [ + "transformer.layers.3.attn.k_norm.text.weight" + ] + + def test_feed_forward_per_modality(self): + assert map_key( + "multimodal_core.layers.0.feed_forward.mlp.0.down_proj.weight", "mot", IMG_TEXT + ) == ["transformer.layers.0.mlp.image.down_proj.weight"] + assert map_key( + "multimodal_core.layers.0.feed_forward.post_attention_layernorm.1.weight", + "mot", + IMG_TEXT, + ) == ["transformer.layers.0.mlp_norm.text.weight"] + + def test_modality_order_swaps_with_index(self): + swapped = {"0": "text", "1": "image"} + assert map_key("multimodal_core.layers.0.self_attn.q_proj.0.weight", "mot", swapped) == [ + "transformer.layers.0.attn.q_proj.text.weight" + ] + + def test_final_norm_fans_out(self): + assert map_key("multimodal_core.norm.weight", "mot", IMG_TEXT) == [ + "transformer.norm.weight", + "transformer.mot_norms.image.weight", + "transformer.mot_norms.text.weight", + ] + + +class TestMapKeyUnsupportedArch: + def test_cross_attention_layer_key_raises(self): + with pytest.raises(NotImplementedError): + map_key("multimodal_core.layers.0.self_attn.q_proj.weight", "cross_attention") + + +# --------------------------------------------------------------------------- +# map_key — JD source -> MoT target (init MoT from a dense JD checkpoint) +# --------------------------------------------------------------------------- + + +class TestMapKeyJDtoMoT: + """Dense JD weights duplicated into both MoT modality copies.""" + + def test_attention_duplicated_into_both_modalities(self): + assert map_key( + "multimodal_core.layers.3.self_attn.q_proj.weight", "joint_decoder", target_arch="mot" + ) == [ + "transformer.layers.3.attn.q_proj.image.weight", + "transformer.layers.3.attn.q_proj.text.weight", + ] + assert map_key( + "multimodal_core.layers.3.self_attn.q_norm.weight", "joint_decoder", target_arch="mot" + ) == [ + "transformer.layers.3.attn.q_norm.image.weight", + "transformer.layers.3.attn.q_norm.text.weight", + ] + + def test_layernorms_duplicated(self): + assert map_key( + "multimodal_core.layers.0.input_layernorm.weight", "joint_decoder", target_arch="mot" + ) == [ + "transformer.layers.0.attn_norm.image.weight", + "transformer.layers.0.attn_norm.text.weight", + ] + assert map_key( + "multimodal_core.layers.0.post_attention_layernorm.weight", + "joint_decoder", + target_arch="mot", + ) == [ + "transformer.layers.0.mlp_norm.image.weight", + "transformer.layers.0.mlp_norm.text.weight", + ] + + def test_mlp_duplicated(self): + assert map_key( + "multimodal_core.layers.1.mlp.down_proj.weight", "joint_decoder", target_arch="mot" + ) == [ + "transformer.layers.1.mlp.image.down_proj.weight", + "transformer.layers.1.mlp.text.down_proj.weight", + ] + + def test_final_norm_fans_out(self): + assert map_key("multimodal_core.norm.weight", "joint_decoder", target_arch="mot") == [ + "transformer.norm.weight", + "transformer.mot_norms.image.weight", + "transformer.mot_norms.text.weight", + ] + + def test_shared_parts_unchanged(self): + # embedding / head / adapter map the same regardless of target arch. + assert map_key("text_preprocessor.embed.weight", "joint_decoder", target_arch="mot") == [ + "transformer.token_embedding.embedding.weight" + ] + assert map_key("adapter.model.ln_q.weight", "joint_decoder", target_arch="mot") == [ + "adapter.ln_q.weight" + ] + + +# --------------------------------------------------------------------------- +# Completeness: mapping covers every target key for a tiny JD / MoT model +# --------------------------------------------------------------------------- + + +def _tiny_config() -> ModelConfig: + return ModelConfig( + dim=64, + n_layers=2, + n_heads=4, + n_kv_heads=2, + head_dim=16, # decoupled (4 * 16 = 64 here, but exercises the field) + qk_norm=True, + ffn_hidden_dim=128, + vocab_size=100, + max_seq_len=64, + ) + + +def _target_keys(model: Transformer, adapter: torch.nn.Module) -> set[str]: + return {f"transformer.{k}" for k in model.state_dict()} | { + f"adapter.{k}" for k in adapter.state_dict() + } + + +def _adapter(): + cfg = AdapterConfig(type="mlp_2layer", pre_norm="rmsnorm", hidden_dim=32, activation="gelu") + return build_adapter(cfg, in_dim=32, out_dim=64) + + +def _jd_source_keys(n_layers: int) -> list[str]: + keys = [ + "text_preprocessor.embed.weight", + "text_head.weight", + "multimodal_core.norm.weight", + "adapter.model.ln_q.weight", + "adapter.model.mlp.0.weight", + "adapter.model.mlp.0.bias", + "adapter.model.mlp.2.weight", + "adapter.model.mlp.2.bias", + ] + for i in range(n_layers): + p = f"multimodal_core.layers.{i}" + keys += [ + f"{p}.self_attn.q_proj.weight", + f"{p}.self_attn.k_proj.weight", + f"{p}.self_attn.v_proj.weight", + f"{p}.self_attn.o_proj.weight", + f"{p}.self_attn.q_norm.weight", + f"{p}.self_attn.k_norm.weight", + f"{p}.input_layernorm.weight", + f"{p}.post_attention_layernorm.weight", + f"{p}.mlp.gate_proj.weight", + f"{p}.mlp.up_proj.weight", + f"{p}.mlp.down_proj.weight", + ] + return keys + + +def _mot_source_keys(n_layers: int) -> list[str]: + keys = [ + "text_preprocessor.embed.weight", + "text_head.weight", + "multimodal_core.norm.weight", + "adapter.model.ln_q.weight", + "adapter.model.mlp.0.weight", + "adapter.model.mlp.0.bias", + "adapter.model.mlp.2.weight", + "adapter.model.mlp.2.bias", + ] + for i in range(n_layers): + p = f"multimodal_core.layers.{i}" + for m in ("0", "1"): + keys += [ + f"{p}.self_attn.input_layer_norm.{m}.weight", + f"{p}.self_attn.q_proj.{m}.weight", + f"{p}.self_attn.k_proj.{m}.weight", + f"{p}.self_attn.v_proj.{m}.weight", + f"{p}.self_attn.o_proj.{m}.weight", + f"{p}.self_attn.q_norm.{m}.weight", + f"{p}.self_attn.k_norm.{m}.weight", + f"{p}.feed_forward.mlp.{m}.gate_proj.weight", + f"{p}.feed_forward.mlp.{m}.up_proj.weight", + f"{p}.feed_forward.mlp.{m}.down_proj.weight", + f"{p}.feed_forward.post_attention_layernorm.{m}.weight", + ] + return keys + + +class TestMappingCompleteness: + def test_joint_decoder_covers_all_target_keys(self): + cfg = _tiny_config() + model = Transformer(cfg, vlm_config=JointDecoderConfig()) + adapter = _adapter() + target = _target_keys(model, adapter) + + src = {k: torch.zeros(1) for k in _jd_source_keys(cfg.n_layers)} + converted, unmapped = map_state_dict(src, "joint_decoder") + + assert unmapped == [] + assert set(converted) == target # no missing, no unexpected + + def test_mot_covers_all_target_keys(self): + cfg = _tiny_config() + model = Transformer(cfg, vlm_config=MoTConfig(), num_image_tokens=8) + adapter = _adapter() + target = _target_keys(model, adapter) + + src = {k: torch.zeros(1) for k in _mot_source_keys(cfg.n_layers)} + converted, unmapped = map_state_dict(src, "mot", IMG_TEXT) + + assert unmapped == [] + assert set(converted) == target + + def test_jd_to_mot_covers_all_target_keys(self): + # Init MoT from a dense JD checkpoint: duplicate each dense weight into + # both modality copies. The JD source must fully seed the MoT target. + cfg = _tiny_config() + model = Transformer(cfg, vlm_config=MoTConfig(), num_image_tokens=8) + adapter = _adapter() + target = _target_keys(model, adapter) + + src = {k: torch.zeros(1) for k in _jd_source_keys(cfg.n_layers)} + converted, unmapped = map_state_dict(src, "joint_decoder", target_arch="mot") + + assert unmapped == [] + assert set(converted) == target + + def test_jd_to_mot_symmetric_duplication_shares_source_tensor(self): + # Symmetric init: the image and text copies are the SAME source tensor. + src = {"multimodal_core.layers.0.self_attn.o_proj.weight": torch.zeros(3)} + converted, _ = map_state_dict(src, "joint_decoder", target_arch="mot") + assert ( + converted["transformer.layers.0.attn.o_proj.image.weight"] + is converted["transformer.layers.0.attn.o_proj.text.weight"] + )