Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add gradient checkpointing for DINOv3 and DGTRS-CLIP (bs 8→24)

Browse Source 
- DINOv3: checkpoint each of 24 transformer blocks (recompute on backward)
- DGTRS-CLIP: checkpoint each of 12 transformer blocks
- Enables batch_size=24 on RTX 4090 (was 8 without checkpointing)
- Peak VRAM: 20.3 GB at bs=24 (was OOM at bs=16 before)
- ~20-30% slower per step, but 3x more in-batch negatives (23 vs 7)
- Enabled by default (gradient_checkpointing=True in config)
- Update README with VRAM benchmarks and checkpointing docs

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 21:34:31 +03:00
parent da2d2ea90e
commit 6b7bcae198
5 changed files with 56 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
README.md
View File

@@ -186,10 +186,29 @@ Symmetric InfoNCE with learnable temperature (CLIP-style `logit_scale`):
Loss and adapters run in **fp32** (AMP autocast disabled) to prevent gradient overflow.
### Gradient checkpointing
Gradient checkpointing trades compute for VRAM by recomputing activations during
backward instead of storing them. Enabled by default (`gradient_checkpointing=True`).
| Component | Without checkpointing | With checkpointing |
|-----------|:---:|:---:|
| DINOv3 (24 blocks) | stores all 24 block activations | recomputes on backward |
| DGTRS-CLIP (12 blocks) | stores all 12 block activations | recomputes on backward |
| **Max batch_size** (RTX 4090, shared encoder) | **8** | **24** |
| Speed penalty | — | ~20-30% slower per step |
VRAM tested on RTX 4090 (24 GB) with shared DINOv3 WEB + DGTRS-CLIP + text:
| `batch_size` | Peak VRAM | Status |
|:---:|:---:|:---:|
| 16 | 14.7 GB | OK |
| 24 | 20.3 GB | OK (recommended) |
| 32 | >24 GB | OOM |
### Gradient accumulation
With `batch_size=8` on a 24 GB GPU, VRAM is the bottleneck. Gradient accumulation
emulates a larger effective batch without extra memory:
Gradient accumulation emulates a larger effective batch without extra memory:
```
effective_batch_size = batch_size × grad_accum_steps
@@ -197,17 +216,17 @@ effective_batch_size = batch_size × grad_accum_steps
| Setting | `batch_size` | `grad_accum_steps` | Effective batch | In-batch negatives |
|---------|:---:|:---:|:---:|:---:|
| Default | 8 | 1 | 8 | 7 |
| Recommended | 8 | 8 | 64 | 7 per micro-batch |
| Default | 24 | 1 | 24 | 23 |
| Large effective batch | 24 | 4 | 96 | 23 per micro-batch |
**Note:** gradient accumulation averages gradients across micro-batches, but each
micro-batch still only sees `batch_size` in-batch negatives. To increase the number
of negatives per forward pass, increase `batch_size` directly (requires more VRAM).
```bash
# Example: effective batch of 64 with 8 accumulation steps
# Example: effective batch of 96 with gradient accumulation
python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
--filter-meta meta/seg_filter.json --grad-accum 8
--filter-meta meta/seg_filter.json --batch-size 24 --grad-accum 4
```
### Metrics