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

@@ -74,6 +74,7 @@ class TrainConfigGTAUAV:
init_gate: float = 0.7
baseline_mode: bool = False
shared_encoder: bool = True # single DINOv3 WEB for both branches (saves ~4-5 GB)
gradient_checkpointing: bool = True # trade compute for VRAM (allows larger batch)
# Training.
resume_from: str | None = None # path to checkpoint for resuming
@@ -354,6 +355,17 @@ def train(cfg: TrainConfigGTAUAV) -> None:
device=cfg.device,
).to(cfg.device)
# --- Gradient checkpointing (trade compute for VRAM) ---
if cfg.gradient_checkpointing:
if cfg.shared_encoder:
model.image_encoder.set_gradient_checkpointing(True)
else:
model.drone_encoder.set_gradient_checkpointing(True)
model.sat_encoder.set_gradient_checkpointing(True)
if model.text_encoder is not None:
model.text_encoder.transformer.gradient_checkpointing = True
LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)")
n_trainable = sum(p.numel() for p in model.trainable_parameters())
n_total = sum(p.numel() for p in model.parameters())
LOGGER.info(