startrail-org--pixelrag/train/docs/training_dev_notes.md

Training Developer Notes

This file keeps internal training notes that are useful for maintainers but too noisy for the top-level reproduction README.

Core Files

File	Description
train_contrastors.py	Primary training script for the reproduced embedding fine-tune. Uses InfoNCE, hard negatives, GradCache, and optional ViT LoRA.
mine_hard_negatives.py	Mines near-miss documents with the base model for hard-negative training.
filter_hard_negatives_vqa.py	Filters mined retrieval candidates with a VLM so false negatives are not used as hard negatives.
run_filter_hard_negatives_chunks.py	Runs VQA hard-negative filtering in fixed-size chunks, useful for large JSONL files.
verify_embeddings.py	Compares base vs fine-tuned embeddings with similarity margins and retrieval metrics.
tests/test_grad_equivalence.py	Single-GPU gradient correctness tests for GradCache math.
tests/test_grad_multi_gpu.py	Multi-GPU DDP/gather/loss-scaling gradient correctness tests.

Legacy / secondary scripts:

• train_colpali.py — HF Trainer-based path. Simpler, but not the reproduced training run.
• train_swift.py — ms-swift alternative.
• train.py, model.py, dataset.py, evaluate.py — older local training/eval code.

Data Formats

Basic JSONL format:

{"query": "What is the population of Tokyo?", "chunk_path": "/path/to/chunk.png"}

Hard-negative format:

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

Hard-negative-with-retrieval-candidates format, used before VQA filtering:

{"query": "...", "chunk_path": "...", "neg_chunk_paths": ["..."], "retrieve_top20": [{"rank": 1, "path": "...", "score": 0.61}]}

Notes:

• chunk_path is resolved relative to the JSONL file when paths are relative.
• Training images are screenshot chunks; last chunks can be smaller than the common tile size.
• neg_chunk_paths should contain mined hard negatives, not random negatives.

Training Pipeline

1. Data loading pre-validates images at init so all DDP ranks process the same number of batches.
2. BiQwen3Processor handles text tokenization and image preprocessing with visual-token resolution control.
3. BiQwen3 embeds the text query and image document into single L2-normalized vectors.
4. InfoNCE is computed over the similarity matrix.
5. With hard negatives, docs are interleaved as [pos, neg1, neg2, pos, neg1, neg2, ...].
6. Multi-GPU training uses gather_with_grad so document embeddings from other ranks contribute gradients.
7. GradCache keeps activation memory tied to --grad-cache-chunk rather than the full effective batch.

VQA Filtering For Hard Negatives

filter_hard_negatives_vqa.py removes false negatives from mined retrieval candidates:

1. Read retrieve_top20.
2. Skip the positive chunk_path.
3. Check up to the first K non-positive candidates (--candidate-k, default 10).
4. Ask the VLM to answer the query from each candidate image.
5. Judge that answer on the same image.
6. If verdict is CORRECT, treat the candidate as a false negative and skip it.
7. If verdict is WRONG or CANNOT_ANSWER, keep it as a hard negative.
8. Stop after collecting --num-hard-negatives hard negatives.
9. Skip the example if not enough hard negatives are found within the first K candidates.

Example:

OPENAI_API_KEY=... python filter_hard_negatives_vqa.py \
    --input /tmp/sample_100_hn.jsonl \
    --output /tmp/sample_100_hn_v2.jsonl \
    --reviews-output /tmp/sample_100_hn_v2.reviews.jsonl \
    --summary-output /tmp/sample_100_hn_v2.summary.json \
    --candidate-k 10 \
    --num-hard-negatives 2 \
    --concurrency 8

For large files:

OPENAI_API_KEY=... python run_filter_hard_negatives_chunks.py \
    --input training/data/lite-query-v2-full-filtered-hn.jsonl \
    --output-dir training/data/lite-query-v2-full-filtered-hn-v2-chunks \
    --chunk-size 10000 \
    --candidate-k 10 \
    --num-hard-negatives 2 \
    --concurrency 8 \
    --skip-existing

Each chunk folder contains:

• filtered_hn.jsonl
• candidate_reviews.jsonl
• summary.json

summary.json is updated incrementally and tracks missing-path ratios for positive paths, reviewed candidate paths, and all checked paths combined.

Dataset Packaging

For regenerating or uploading a Hugging Face dataset:

# Prepare HF dataset folder: convert absolute paths to relative paths and hardlink images.
python prepare_hf_dataset.py \
    --split-dir training/data/lite-query-v2-full-filtered-hn-v2-chunks/split \
    --image-root /opt/dlami/nvme/kiwix_tiles \
    --output-dir hf_dataset_export/screenshot-training

# Package images into tar shards.
python package_hf_image_shards.py \
    --source-dir hf_dataset_export/screenshot-training \
    --output-dir hf_dataset_export_sharded/screenshot-training

# Upload to Hugging Face.
python upload_hf_dataset.py \
    --local-dir hf_dataset_export_sharded/screenshot-training

Upload requires a Hugging Face token with write permission.

Test Commands

# Single-GPU: GradCache math + RandContext + clip_loss + rope_deltas
CUDA_VISIBLE_DEVICES=0 uv run python tests/test_grad_equivalence.py

# Multi-GPU: DDP + gather + loss scaling + gradient sync
CUDA_VISIBLE_DEVICES=0,1 uv run torchrun --nproc_per_node=2 tests/test_grad_multi_gpu.py

Archived / Experimental Mode

train_contrastors.py still has a query-side-tune mode that trains only the query tower while keeping the doc/image tower frozen, so datastore embeddings do not change. This is not part of the reproduction path in README.md.

Important caveats:

• Query-side retrieval eval depends on an external search API that accepts pre-computed embedding queries.
• --query-side-backward direct was the stable path in smoke tests.
• --query-side-backward gradcache was experimental and may hang in multi-GPU real-data runs.