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

GradCache Training (train_contrastors.py)

GradCache training for Qwen3-VL-Embedding, adapted from nomic-ai/contrastors. Verified correct via gradient equivalence tests (single-GPU + multi-GPU).

Scope Note

This document describes the original GradCache path used by train_contrastors.py in standard mode.

For --mode query-side-tune:

• default backward path is --query-side-backward direct
• --query-side-backward gradcache still exists, but is currently experimental
• multi-GPU real-data smoke tests were stable with direct and unstable with gradcache

Why GradCache?

Contrastive learning benefits from large batch sizes (more negatives = better). But GPU memory limits batch size. GradCache breaks this constraint:

Standard:  batch=4, chunk=4 → memory of 4, negatives = 4
GradCache: batch=4, chunk=2 → memory of 2, negatives = 4
Multi-GPU: batch=4, chunk=2, 5 GPUs → memory of 2, negatives = 20

How It Works

GradCache splits each batch into small chunks and processes them in 3 steps:

1. Forward all chunks WITHOUT grad → cache embeddings + RNG states (constant memory)
2. Compute InfoNCE loss on ALL cached embeddings → get embedding gradients via backward on detached tensors
3. Replay forward WITH grad using saved RNG states, apply surrogate loss (dot(emb, cached_grad)) → real parameter gradients

By the chain rule, d(loss)/d(θ) = d(loss)/d(emb) · d(emb)/d(θ). Step 2 computes the first factor, step 3 computes the second. The result is mathematically identical to a full-batch backward pass.

DDP Integration

Multi-GPU adds two distributed primitives:

• gather_with_grad: All-gathers document embeddings across ranks (step 2). Backward does reduce_scatter to distribute gradients back.
• Manual all_reduce(AVG): All surrogate backward calls run under no_sync() to avoid DDP reducer deadlocks (query chunks skip the visual encoder while doc chunks use it → different "used" parameter sets confuse find_unused_parameters). Gradients are manually synced after all chunks.

The loss is scaled by world_size before backward, so after all_reduce(AVG):

final_grad = (1/W) × Σ_r [W × d(CE_r)/d(θ)] = Σ_r d(CE_r)/d(θ) = d(total_CE)/d(θ)

Key Components

Component	Source	Purpose
RandContext	contrastors/rand_state.py	Save/restore GPU RNG state for dropout replay
gather_with_grad	contrastors/distributed.py	All-gather with gradient flow (backward = reduce_scatter)
clip_loss	contrastors/loss.py	InfoNCE with learnable logit scale + hard negative support
grad_cache_loss	contrastors/loss.py	Full GradCache pipeline (3-step)
LogitScale	contrastors/OpenCLIP	Learnable log_scale parameter, clamped post-step
BiQwen3	colpali-engine	Qwen3-VL wrapped as bi-encoder (last-token pool + L2 norm)

LogitScale

Learnable temperature in log-space, initialized to ln(1/0.07) ≈ 2.66:

forward:  similarity * exp(log_scale)     # no clamp in forward (avoids gradient dead zone)
after optimizer.step():  log_scale.clamp_(0, ln(100))  # contrastors pattern

Quick Start

PYTHON=.venv-sglang/bin/python

# Single GPU
CUDA_VISIBLE_DEVICES=3 $PYTHON training/train_contrastors.py \
    --max-steps 500 --batch-size 4 --grad-cache-chunk 2

# Multi-GPU (5 GPUs, cross-GPU negatives + GradCache)
CUDA_VISIBLE_DEVICES=3,4,5,6,7 .venv-sglang/bin/torchrun --nproc_per_node=5 \
    training/train_contrastors.py --max-steps 500 --batch-size 8 --grad-cache-chunk 2

# With hard negatives (requires train_hn.jsonl from mine_hard_negatives.py)
CUDA_VISIBLE_DEVICES=3,4,5,6,7 .venv-sglang/bin/torchrun --nproc_per_node=5 \
    training/train_contrastors.py --train-jsonl training/data/train_hn.jsonl \
    --num-hard-negatives 2 --batch-size 4 --grad-cache-chunk 2

# Resume from checkpoint
$PYTHON training/train_contrastors.py \
    --resume training/output_contrastors/checkpoint-200

Hyperparameters

Parameter	Default	Notes
--batch-size	4	Per-GPU batch size
--grad-cache-chunk	2	Chunk size for GradCache (memory = this × per-sample cost)
--lr	2e-5	Peak learning rate (cosine schedule)
--warmup-steps	50	Linear warmup
--max-steps	500	Total training steps
--temperature	0.07	Initial temperature (learnable logit scale = 1/temp)
--num-hard-negatives	0	Hard negs per query (docs interleaved: [pos, neg1, neg2, ...])
--lora-r	32	LoRA rank
--lora-alpha	32	LoRA alpha
--max-num-visual-tokens	256	Image resolution (~200K pixels)
--max-grad-norm	1.0	Gradient clipping (model + logit_scale)

Comparison with train_colpali.py

Feature	train_colpali.py	train_contrastors.py
Training infra	HF Trainer (ContrastiveTrainer)	Custom loop
GradCache	No	Yes
Cross-GPU negatives	all_gather	gather_with_grad
Temperature	Fixed	Learnable (LogitScale)
Hard negatives	No	Yes (--num-hard-negatives)
DDP strategy	Standard DDP	no_sync + manual all_reduce
Checkpoint/resume	HF Trainer	Manual

Verified Correct

Gradient equivalence tests confirm GradCache produces identical gradients to a full-memory reference (15 tests, all passing):

Single-GPU (tests/test_grad_equivalence.py):

• GradCache chain-rule decomposition: cosine ≥ 0.9999 for chunk_size = 1, 2, batch_size
• RandContext dropout replay: cosine ≥ 0.9999 for all chunk sizes
• clip_loss label arithmetic: basic, hard negatives, divisibility assertion
• _clear_rope_deltas: prevents image→text rope state leakage

Multi-GPU (tests/test_grad_multi_gpu.py, 2×GPU):

• GradCache DDP vs reference: cosine ≥ 0.9998 (with and without dropout)
• gather_with_grad backward: reduce_scatter gives correct gradient = W
• loss*W + all_reduce(AVG) = gradient of total loss: exact match
• Gradients identical across ranks after sync: max diff = 0

# Run tests
CUDA_VISIBLE_DEVICES=2 python training/tests/test_grad_equivalence.py
CUDA_VISIBLE_DEVICES=2,3 torchrun --nproc_per_node=2 training/tests/test_grad_multi_gpu.py

Code Attribution

Core GradCache implementation adapted from nomic-ai/contrastors:

• contrastors/loss.py — grad_cache_loss, clip_loss, get_chunked_embeddings
• contrastors/rand_state.py — RandContext
• contrastors/distributed.py — gather_with_grad
• contrastors/models/biencoder/modeling_biencoder.py — LogitScale
• contrastors/trainers/text_text.py — post-step clamp_() pattern