startrail-org--pixelrag/train/CLAUDE.md

CLAUDE.md

⚠️ Dev notes may be out of date. docs/training_dev_notes.md (and some references in this file) describe historical internals and may have drifted from the current code — e.g. they mention train_colpali.py and an old training.* package/trainer stack (train.py, evaluate.py, dataset.py, model.py) that are superseded by the self-contained train_contrastors.py. Treat the docs as background; trust the code and README.md as source of truth.

Pinned Versions

All training and eval must use these exact versions — mismatches cause silent numerical divergence:

Package	Version
PyTorch	2.9.1+cu129
cuDNN	92000
transformers	4.57.1

uv sync will install the correct versions from pyproject.toml + lockfile. cuDNN 9.20 is forced via override-dependencies in pyproject.toml (torch 2.9.1 ships with 9.10, but we need 9.20 for native bf16 Conv3d).

Always use uv run to ensure the locked environment is used.

Training Pipeline (Best Config)

# 0. Install
uv sync

# 1. Mine hard negatives (requires search API at localhost:30888)
uv run python mine_hard_negatives.py \
    --input training/data/train.jsonl \
    --output training/data/train_hn.jsonl \
    --num-negatives 7 --n-docs 50 \
    --filter-mode margin --margin 0.95

# 2. Train (best config)
CUDA_VISIBLE_DEVICES=1,2 uv run torchrun --nproc_per_node=2 train_contrastors.py \
    --data-split-dir training/data/lite-query-v2-full-filtered-hn-v2-chunks/split \
    --max-steps 50 \
    --batch-size 16 \
    --grad-cache-chunk 4 \
    --num-hard-negatives 5 \
    --lr 1e-5 \
    --warmup-steps 20 \
    --test-eval-steps 50 \
    --test-max-pairs 0 \
    --eval-steps 25 \
    --save-steps 50 \
    --output-dir training/output_nvme/output_hn_best

# 3. Eval checkpoint on hard-mini-v6 (default) or v7
cd /home/user/Vis-RAG/agent && \
CUDA_VISIBLE_DEVICES=4 uv run python scripts/eval_lora_checkpoint.py \
    /home/user/wiki-screenshot-training/training/output_nvme/output_hn_best/checkpoint-50 \
    --tiles-dir tiles-hard-mini-v6

# Eval on hard-mini-v7 (400 queries, 7723 tiles — more comprehensive)
# ⚠️ v7 has known issues — prefer v6 for now
cd /home/user/Vis-RAG/agent && \
CUDA_VISIBLE_DEVICES=4 uv run python scripts/eval_lora_checkpoint.py \
    /home/user/wiki-screenshot-training/training/output_nvme/output_hn_best/checkpoint-50 \
    --tiles-dir tiles-hard-mini-v7 \
    --vllm-url http://localhost:8201/v1 \
    --vllm-model Qwen/Qwen3-VL-4B-Instruct

# Retrieval-only (skip QA, faster)
cd /home/user/Vis-RAG/agent && \
CUDA_VISIBLE_DEVICES=4 uv run python scripts/eval_lora_checkpoint.py \
    /path/to/checkpoint --tiles-dir tiles-hard-mini-v7 --retrieval-only

# Resume from checkpoint
uv run python train_contrastors.py --resume training/output_nvme/output_hn_best/checkpoint-50

Query-Side Tune

train_contrastors.py also supports a query-only fine-tuning mode:

• --mode query-side-tune trains the query tower only
• the doc/image tower stays frozen at the base model
• datastore embeddings therefore stay valid across checkpoints
• checkpoint eval can query the full external datastore directly via wiki-screenshot search API

Recommended command:

OPENAI_API_KEY=<your-key> \
OPENAI_BASE_URL=https://us.api.openai.com/v1 \
CUDA_VISIBLE_DEVICES=2,3,4,5 uv run torchrun --master_port 29531 --nproc_per_node=4 train_contrastors.py \
    --mode query-side-tune \
    --query-side-backward direct \
    --data-split-dir training/data/lite-query-v2-full-filtered-hn-v2-chunks/split \
    --batch-size 64 \
    --num-hard-negatives 2 \
    --lr 2e-6 \
    --warmup-steps 20 \
    --eval-steps 25 \
    --test-eval-steps 50 \
    --save-steps 100 \
    --max-steps 100 \
    --max-num-visual-tokens 1024 \
    --search-api-url http://localhost:30888 \
    --simpleqa-max-examples 100 \
    --vllm-url http://localhost:8201/v1 \
    --vllm-model Qwen/Qwen3-VL-4B-Instruct \
    --output-dir training/output_nvme/output_query_side

Required services:

• search API on :30888 — for retrieval eval (recall@1/3)
• vLLM on :8201 — Qwen3-VL-4B-Instruct for VQA answering
• OpenAI API — GPT-4.1 as SimpleQA grader (needs OPENAI_API_KEY)

Important defaults / caveats:

• In query-side-tune, default --query-side-backward is direct (no GradCache replay).
• This is intentional: direct was stable in 4-GPU real-data smoke tests.
• --query-side-backward gradcache is currently experimental for query-side mode and may hang.
• standard mode is unchanged and still uses the original GradCache path.
• Query-side retrieval eval requires a new enough wiki-screenshot search_api.py that accepts pre-computed embedding queries.

Training Scripts

Script	Description
train_contrastors.py	Primary. Standard mode uses GradCache; query-side mode freezes the doc tower and defaults to direct backward.
train_swift.py	ms-swift alternative. Simpler (one sft_main() call) but no GradCache / learnable temp.
train_colpali.py	Legacy HF Trainer-based. Simpler but no GradCache.

Key args for train_contrastors.py:

Arg	Default	Notes
--mode	standard	standard = shared tower training; query-side-tune = train query side only
--query-side-backward	direct	Only used in query-side-tune; direct is stable default
--batch-size	4	Per-GPU. Effective negatives = batch × num_gpus
--grad-cache-chunk	2	Memory ∝ this, not batch size
--lr	2e-5	Cosine schedule with warmup
--warmup-steps	50	Linear warmup
--temperature	0.07	Initial temp (learnable via LogitScale)
--num-hard-negatives	0	Set 2–7 when using train_hn.jsonl
--lora-r / --lora-alpha	32 / 32	LoRA rank and alpha
--lora-vit	off	Also apply LoRA to ViT vision encoder (see below)
--max-num-visual-tokens	1024	Image resolution control
--max-steps	500	Total training steps
--eval-steps	100	Validation loss frequency
--test-eval-steps	250	Full retrieval eval (R@1/5/10, MRR)
--text-warmup-steps	0	Text-only warmup steps before image training
--text-data-dir	None	Directory with text-qa-pair JSONL files
--text-mix-ratio	0	Fraction of text batches during image phase
--text-curriculum	off	Gradual text→image transition (50%→33%→20%→0%)
--test-batch-size	16	Eval batch size (lower to avoid OOM on v7/v8)
--hardness-alpha	0	LLaVE hardness weighting (0=off, try 5–9). Upweights harder negatives in softmax

LoRA Target Modules

By default, LoRA only targets the LLM backbone attention layers (q_proj, k_proj, v_proj, o_proj). The ViT vision encoder is not tuned — its attention uses fused qkv naming that doesn't match.

Qwen3-VL architecture:

ViT (frozen by default):
  model.visual.blocks.N.attn.qkv        — fused QKV (Linear)
  model.visual.blocks.N.attn.proj        — output proj (Linear)
  model.visual.blocks.N.mlp.linear_fc1   — MLP up (Linear)
  model.visual.blocks.N.mlp.linear_fc2   — MLP down (Linear)

Merger (vision→LLM bridge):
  model.visual.merger.linear_fc1/fc2     — not matched by --lora-vit

LLM backbone (tuned by default):
  model.layers.N.self_attn.q_proj/k_proj/v_proj/o_proj

With --lora-vit, LoRA is also applied to ViT attn + MLP layers:

• Without: 224 LoRA layers, ~12.8M trainable params
• With --lora-vit: 416 LoRA layers, ~25M trainable params

This matches the colpali/BiQwen2 approach which uses .*model.* regex to target all layers including ViT.

Query-side-specific eval behavior:

• test split: local query embeddings are sent to --search-api-url, retrieve top-3, report recall@1 / recall@3
• SimpleQA: current query tower retrieves top-3 from the same endpoint, reports article-level recall@1 / recall@3, then grades answers with an OpenAI-compatible judge

Data

Source: LLM-Augmented Query-Document Pairs

Training pairs are generated by sending Wikipedia screenshot chunks to LLMs, which produce natural-language queries that a user might ask to find that specific chunk.

Batch JSONL files live in agent/scripts/contrastive/batches/:

Batch range	Model	Pairs	Notes
batch_000.jsonl – batch_057.jsonl	Gemini 3.1 Pro	40,402	Higher quality, more diverse queries
batch_200.jsonl – batch_257.jsonl	Flash-Lite	88,418	Cheaper, noisier, larger volume

Total: ~128,820 augmented pairs across 116 batch files.

Data Format

Basic format (one JSON object per line):

{"query": "What is the population of Tokyo?", "chunk_path": "/opt/dlami/nvme/kiwix_tiles/shard_000/shard_00042/350170.png.tiles/chunk_0000_00.png"}

Hard negative format (after mining):

{"query": "...", "chunk_path": "...", "neg_chunk_paths": ["/path/to/neg1.png", "/path/to/neg2.png"]}

Train / Val / Test Split

Split the combined batch data before training:

1. Concatenate all batch JSONL files into one pool
2. Shuffle with a fixed seed for reproducibility
3. Split by ratio:
   • Train: 90% — used for contrastive fine-tuning
   • Val (eval.jsonl): 5% — monitored every --eval-steps for loss
   • Test (test.jsonl): 5% — full retrieval eval (R@1, R@5, R@10, MRR) every --test-eval-steps
4. Mine hard negatives on the train split only:

   uv run python mine_hard_negatives.py \
       --input training/data/train.jsonl \
       --output training/data/train_hn.jsonl \
       --num-negatives 7 --n-docs 20 --margin 0.95
   

Output files in training/data/:

• train.jsonl / train_hn.jsonl — training (with optional hard negs)
• eval.jsonl — validation
• test.jsonl — held-out test

External Training Datasets (MOCA / MADQA)

Additional contrastive training data from MOCA and MADQA benchmarks, with pre-mined hard negatives:

Dataset	HF Link	Rows	Images	Upload Status
MOCA ColPali	Chrisyichuan/moca-colpali-training	118,195	118,195	✅ done
MOCA PixelRAG Ind	Chrisyichuan/moca-visrag-ind-training	122,752	122,752	✅ done
MOCA PixelRAG Syn	Chrisyichuan/moca-visrag-syn-training	239,206	239,298	✅ done
MADQA	Chrisyichuan/madqa-training	1,840	3,598	✅ done

Images are stored as tar shards under image_shards/. After downloading, extract with the included script:

# Download all four
pip install huggingface_hub
for repo in moca-colpali-training moca-visrag-ind-training moca-visrag-syn-training madqa-training; do
    huggingface-cli download Chrisyichuan/$repo --repo-type dataset --local-dir /opt/dlami/nvme/external_data/$repo
done

# Extract images from tar shards
cd /opt/dlami/nvme/external_data/<dataset>
python extract_hf_image_shards.py --dataset-dir .

Each dataset contains a JSONL metadata file + image_shards/ directory. After extraction, images/ is created. Format matches the main training data:

{"query": "...", "chunk_path": "images/...", "neg_chunk_paths": ["images/...", "..."], "source_dataset": "moca"}

Preparation scripts: prepare_andy_datasets.py (build HF folders from raw JSONL), package_andy_shards.py (tar shard packaging), upload_andy_datasets.py (upload to Hub).

Generating Synthetic Data (for quick testing)

fake_data.py generates template-based queries from article titles (not LLM-augmented):

uv run python fake_data.py \
    --tiles-dir /opt/dlami/nvme/kiwix_tiles \
    --articles-json /opt/dlami/nvme/kiwix/wikipedia_en_all_maxi_2025-08.zim.articles.json \
    --output-dir training/data --num-articles 1000

Evaluation

# Verify fine-tuned vs base embeddings
CUDA_VISIBLE_DEVICES=0 uv run python verify_embeddings.py \
    --adapter training/output_nvme/output_contrastors/checkpoint-500 --max-pairs 100

# Gradient correctness tests
CUDA_VISIBLE_DEVICES=0 uv run python tests/test_grad_equivalence.py
CUDA_VISIBLE_DEVICES=0,1 uv run torchrun --nproc_per_node=2 tests/test_grad_multi_gpu.py

Checkpoint Hygiene

Checkpoints produced during training quickly fill up the disk. Periodically clean up checkpoints that are not promising:

• After each new checkpoint is produced, check whether training/output_nvme/ contains old checkpoints that are clearly unneeded (e.g. intermediate versions where loss did not drop or R@1 did not improve).
• If you are sure they are not promising, delete them directly.
• If you are unsure whether to delete, ask the user before acting.

Launching Training

⚠️ All training jobs must be launched inside a tmux session! Running training in a bare terminal = everything is lost when the SSH connection drops. We have already lost a training run once because of this. Never run torchrun directly in a bare SSH terminal.

# Correct approach: launch inside tmux
tmux new-session -d -s train -c /home/ubuntu/wiki-screenshot-training
tmux send-keys -t train "CUDA_VISIBLE_DEVICES=1,2,3 uv run torchrun --nproc_per_node=3 train_contrastors.py \
    [args...] 2>&1 | tee training/output_nvme/<run>/train.log" Enter

# View training output
tmux attach -t train

Never pipe training output through head/tail/etc. — SIGPIPE will kill the torchrun workers.

⚠️ Before launching training you must confirm that OPENAI_API_KEY and OPENAI_BASE_URL are set! Without these two environment variables, the QA score of the mini-v6 eval will silently return 0 (the grader swallows all exceptions), and you will think the model is bad when in fact grading never ran. If you find that the current shell has no OPENAI_API_KEY, you must immediately remind the user to set it and not continue launching training.

These two variables should already be configured in ~/.zshrc. Verify before launching:

# Mandatory check before launching training
echo "OPENAI_API_KEY=${OPENAI_API_KEY:+SET}" "OPENAI_BASE_URL=${OPENAI_BASE_URL}"

This API key requires us.api.openai.com — using the default api.openai.com will 401 silently (the grader catches all exceptions, resulting in QA score = 0 with no visible error).

export OPENAI_API_KEY="sk-proj-..."
export OPENAI_BASE_URL="https://us.api.openai.com/v1"

Experiment Tracking (CSV)

Every training run must write its experiment results to a CSV file for easy side-by-side comparison.

File location: training/output_nvme/<run_name>/metrics.csv

CSV format:

step,eval_loss,eval_acc,recall@1,recall@3,qa_score,peak_eval_acc,peak_qa_score
0,,,0.125,0.300,0.42,0.000,0.42
50,1.05,0.62,0.138,0.325,0.45,0.62,0.45
100,0.98,0.65,0.142,0.340,0.48,0.65,0.48

Metrics that must be recorded:

• step: training step count
• eval_loss: loss on the eval split
• eval_acc: accuracy on the eval split
• recall@1, recall@3: retrieval recall of the test eval
• qa_score: QA score of the test eval (the primary optimization target)
• peak_eval_acc: the highest eval accuracy up to the current step
• peak_qa_score: the highest QA score up to the current step

Each run must also record training/output_nvme/<run_name>/run_config.md:

# Run: <run_name>
- **Ablation**: describe what this experiment is ablating (which baseline it compares against, what variable was changed)
- **Date**: launch date
- **Machine**: machine name (e.g. colin3)
- **GPUs**: the GPU IDs and count used
- **Key args**: list all non-default parameters
- **Baseline**: name of the baseline run being compared against
- **Hypothesis**: expected effect

How to read: After training finishes, or during a mid-run check, parse step-level metrics from train.log and write them to the CSV. If the CSV already exists, append new rows.

Key Findings

• Hard negatives are critical for meaningful improvement beyond baseline
• Primary optimization target is QA score, not recall@k. Recall can drop while QA score improves (query embeddings become more semantically useful even if exact chunk match rate falls).

v8r Ablation Results (2026-04-22)

Full stairstep ablation on training/data/natrual_filtered_v2/split (350 steps, bs=64, lr=7e-6, lora-vit, visual_tokens=4096), evaluated on both miniv6 (200q, 5291 tiles) and miniv8 (400q, 7426 tiles). vLLM reader: Qwen3-VL-4B-Instruct. Grader: gpt-4.1-2025-04-14.

Peak metrics across all eval steps

Run	Config	v6 R@1	v6 R@3	v6 QA	v8 R@1	v8 R@3	v8 QA
base	no training	0.650	0.800	0.665	0.688	0.833	0.730
ab1	+ in-batch only	0.720	0.840	0.705	0.750	0.868	0.750
ab2	+ hard negatives	0.715	0.865	0.735	0.748	0.878	0.778
ab3	+ text warmup	0.730	0.855	0.755	0.755	0.893	0.7825
ab4	+ unfreeze ViT	0.730	0.860	0.755	0.760	0.888	0.7825

Final (last checkpoint, step 350) metrics

Run	Final v6 QA	Final v8 QA	Notes
ab1	0.705	0.7225	v8 peaked @step50 (0.750), degraded to 0.7225 — overfit (acc→1.0)
ab2	0.735	0.7425	v8 peaked @step50 (0.778), degraded to 0.7425
ab3	0.745	0.780	close to peak
ab4	0.755	0.7825	== peak, most stable

Key observations

• Perfect QA stairstep (peak v6: 0.665→0.705→0.735→0.755=0.755; peak v8: 0.730→0.750→0.778→0.7825=0.7825)
• ab3 ≈ ab4 in peak — unfreeze ViT did not add peak QA under this budget (350 steps, single GPU). But ab4 final is closer to peak than ab3, so ViT unfreeze adds training stability, not ceiling
• R@1 is not monotone: ab2 (0.715) < ab1 (0.720) on v6 — hard neg trades some R@1 for R@3 / QA
• ab1 / ab2 overfit late: peak QA reached @step50, then degraded. Lesson: for these configs, shorter training (or early stopping) would have landed better final numbers

Run dirs

• training/output_nvme/v8r_base/ (no wandb, eval-only)
• training/output_nvme/v8r_ab1_inbatch/
• training/output_nvme/v8r_ab2_hn2/
• training/output_nvme/v8r_ab3_hn2_tw50/
• training/output_nvme/v8r_ab4_full/

Tile caches live next to test images (test_miniv6/tiles/.tile_cache_*.pt, test_miniv8/tiles/.tile_cache_*.pt). ⚠️ Do not launch 4 runs in parallel from cold — they race to write the cache. Warm the cache first with one eval run, then parallel launches are safe.

Reader SFT (LlamaFactory)

SFT training of Qwen3-VL-4B to do QA on compressed images (the "reader" model).

Separate venv: LlamaFactory has its own dependency environment, do not use the main project's .venv.

cd sft/LlamaFactory
source .venv/bin/activate   # separate venv, already has LlamaFactory + deepspeed + wandb installed

Key dependencies (already installed):

• flash-linear-attention + causal-conv1d (must be installed, otherwise it silently falls back to slow attention)
• torch 2.9.1+cu129 (same as the main project, cuDNN 9.20 to work around the Conv3D fallback issue)
• deepspeed, wandb

Data preparation (compress images + convert to ShareGPT format):

# Run inside the main project venv
uv run python3 sft/prepare_sft_data.py \
    --dataset-dir /mnt/data/hf_datasets/screenshot-training-natural-filtered-v2 \
    --output-dir /mnt/data/sft_data/compressed_3x \
    --compress-ratio 3 --workers 32

Training (4 GPUs, launch inside tmux):

cd sft/LlamaFactory && source .venv/bin/activate
CUDA_VISIBLE_DEVICES=1,2,3,4 \
FORCE_TORCHRUN=1 NNODES=1 NPROC_PER_NODE=4 \
llamafactory-cli train /home/ubuntu/wiki-screenshot-training/sft/train_qwen3vl_compressed.yaml

Key config (sft/train_qwen3vl_compressed.yaml):

• template: qwen3_vl + enable_thinking: false (do not use qwen3_vl_nothink, it is not equivalent)
• LoRA rank 32, lr 1e-5, DeepSpeed ZeRO-2
• W&B project: llamafactory

vLLM Serving

All vLLM instances must use serving/vllm/ — this subproject pins vLLM + transformers + torch via uv.lock.

Locked versions:

• vLLM 0.19.0
• transformers 4.57.6
• torch 2.10.0

cd serving/vllm
uv sync           # first time only
uv run vllm serve Qwen/Qwen3-VL-4B-Instruct \
    --dtype auto --port 8201 --max-model-len 65536 \
    --gpu-memory-utilization 0.8 --api-key dummy

Never use standalone venvs or other users' venvs to serve models. The serving/vllm/uv.lock is the single source of truth for inference-time dependencies.

Port discovery: Before launching training that needs vLLM, check what's already running:

# List active vLLM endpoints
ss -tlnp | grep -E ':8[0-9]{3}\b'
# Verify model at a port
curl -s http://localhost:<port>/v1/models | python3 -m json.tool

Use the actual port in --vllm-url (e.g. http://localhost:8200/v1). Don't assume the default 8201 — it may be on a different port.