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>

README.md

@@ -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

conf/gtauav_balanced.gin

@@ -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

src/models/asymmetric_encoder.py

@@ -169,12 +169,22 @@ class DINOv3ViT(nn.Module):
         ])
         self.norm = nn.LayerNorm(dim)
         self.embed_dim = dim
+        self.gradient_checkpointing = False
+
+    def set_gradient_checkpointing(self, enable: bool = True) -> None:
+        """Enable/disable gradient checkpointing to trade compute for VRAM."""
+        self.gradient_checkpointing = enable
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass. Returns CLS token embedding [B, dim]."""
         x = self.embeddings(x)
         for block in self.layer:
-            x = block(x)
+            if self.gradient_checkpointing and self.training:
+                x = torch.utils.checkpoint.checkpoint(
+                    block, x, use_reentrant=False,
+                )
+            else:
+                x = block(x)
         x = self.norm(x)
         return x[:, 0]  # CLS token
 

src/models/dgtrs/model.py

@@ -77,8 +77,15 @@ class Transformer(nn.Module):
         self.resblocks = nn.Sequential(
             *[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)]
         )
+        self.gradient_checkpointing = False
 
     def forward(self, x: torch.Tensor):
+        if self.gradient_checkpointing and self.training:
+            for block in self.resblocks:
+                x = torch.utils.checkpoint.checkpoint(
+                    block, x, use_reentrant=False,
+                )
+            return x
         return self.resblocks(x)
 
 

src/training/train_gtauav.py

@@ -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(