wehub-resource-sync
148 lines
6.6 KiB
Markdown
148 lines
6.6 KiB
Markdown
# GradCache Training (train_contrastors.py)
|
||
|
||
GradCache training for Qwen3-VL-Embedding, adapted from [nomic-ai/contrastors](https://github.com/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`:
|
||
|
||
```python
|
||
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
|
||
|
||
```bash
|
||
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
|
||
|
||
```bash
|
||
# 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](https://github.com/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
|