From 6b7bcae198ec013b49c1d084bfd9d2f0fcf4dfdc Mon Sep 17 00:00:00 2001 From: pikaliov Date: Tue, 21 Apr 2026 21:34:31 +0300 Subject: [PATCH] =?UTF-8?q?Add=20gradient=20checkpointing=20for=20DINOv3?= =?UTF-8?q?=20and=20DGTRS-CLIP=20(bs=208=E2=86=9224)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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) --- README.md | 31 +++++++++++++++++++++++++------ conf/gtauav_balanced.gin | 1 + src/models/asymmetric_encoder.py | 12 +++++++++++- src/models/dgtrs/model.py | 7 +++++++ src/training/train_gtauav.py | 12 ++++++++++++ 5 files changed, 56 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 8855a39..17087f5 100644 --- a/README.md +++ b/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 diff --git a/conf/gtauav_balanced.gin b/conf/gtauav_balanced.gin index 7fead23..39218c4 100644 --- a/conf/gtauav_balanced.gin +++ b/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 diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index 2bdba23..80073c4 100644 --- a/src/models/asymmetric_encoder.py +++ b/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 diff --git a/src/models/dgtrs/model.py b/src/models/dgtrs/model.py index 50f1c20..2f238fb 100644 --- a/src/models/dgtrs/model.py +++ b/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) diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 2152600..581706e 100644 --- a/src/training/train_gtauav.py +++ b/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(