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

- 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

View File

@@ -26,6 +26,7 @@ TrainConfigGTAUAV.device = "cuda"
TrainConfigGTAUAV.init_gate = 0.7
TrainConfigGTAUAV.baseline_mode = False
TrainConfigGTAUAV.shared_encoder = True
TrainConfigGTAUAV.gradient_checkpointing = True
# ---- Loss ----
TrainConfigGTAUAV.tau_init = 0.07