vision-adapter-probe-checkpoints

Training checkpoints for the vision-adapter probe: a lightweight projector (~25M params) aligning frozen MoonViT-V2 vision embeddings (4096-dim) to a frozen Qwen/Qwen3.5-2B backbone (hidden 2048). Only the projector trains — backbone and vision tower stay frozen.

This is an alignment probe (grokking study), not a released model. See the training repo for the full pipeline.

Training setup

• Optimizer: AdamW, lr 5e-4, warmup 100 steps, grad clip 1.0
• Batch 16, max seq len 4096, bf16 (PRO 6000 96GB) / grad ckpt OFF on big GPUs
• Per-batch adaptive grad checkpointing on B·L² budget (L_MAX=2500, COST_MAX=50M): small batches stay on the fast no-recompute path, worst-bucket batches flip ckpt ON for that step only
• Data: 116k streamed rows (6 n_vis buckets), selective loss on answer+EOS
• Saves every 100 steps + 10-minute time gate; pushed here automatically

Files

File	Contents
projector_step{K}.pt	Full resumable state at step K: projector + optimizer + scaler + monitor + RNG + plan meta + cfg (~150MB)
projector_final_{N}.pt	Final weights + cfg
probe_log.jsonl	Per-step loss/EMA/gnorm/tokens/L/B·L²/ckpt flag + config header
probe_curves.png	Loss curves
runs.jsonl	Run registry entry
train_*.log	Console log

Step checkpoints are resumable: relaunch training with --resume hf --resume-step K (or --resume local with local files).

Usage

import torch
from vision_adapter.core import HourglassProjector

ckpt = torch.load("projector_final_4000.pt", map_location="cpu", weights_only=False)
proj = HourglassProjector(vision_dim=4096, llm_dim=2048)
proj.load_state_dict(ckpt["proj"])
proj.eval()

weights_only=False is required for step ckpts (they embed optimizer, RNG and monitor state). Load only files you trust — i.e. this repo.

Eval status

Heldout-60 alignment set and generation probes are tracked in the training repo (docs/NEXT_STEPS.md). Check probe_log.jsonl for the loss trajectory and grokking window (expected collapse past ~57.6k samples seen).