Qwen3-0.6B VLM example: configs, checkpoint converter, video finetune

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/<run>/<ts>/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/
+```

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 = "<converted CA DCP>"   # 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

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