Model Optimization
RTX 3050 6 GBMemory optimization strategies enabling PaTz training on a 6GB consumer GPU. FP16 + gradient accumulation + frozen backbone reduce peak VRAM to ~4.6 GB.
💻
6 GB
GPU Budget
RTX 3050 VRAM
🧮
4.39M
Trainable Params
PaTz only (ESM-2 frozen)
📦
32
Effective Batch
4 accum × 8 physical
📉
1e-4
Base LR
AdamW, cosine restart
Memory & Training Optimizations
| Strategy | Benefit | Configuration |
|---|---|---|
| ✓FP16 Mixed Precision | ~2× VRAM | torch.cuda.amp.autocast. Activations stored in float16. Master weights in float32. Enabled by default (--no-fp16 flag to disable). |
| ✓Gradient Accumulation (×4) | Effective batch 32 | grad_accum_steps=4, physical batch=8 → effective batch=32. Optimizer step every 4 forward/backward passes. |
| ✓Gradient Clipping | Stability | torch.nn.utils.clip_grad_norm_(params, 1.0). Prevents exploding gradients in Pairformer triangle updates. |
| ✓Frozen ESM-2 Backbone | ~1.2 GB VRAM | ESM-2 650M (1.25B param) backbone frozen. Only 4.39M PaTz parameters are trained. Eliminates gradient storage for backbone. |
| ✓Warmup + Cosine LR | Convergence | warmup_epochs=5 linear warmup → CosineAnnealingWarmRestarts(T_0=10, T_mult=2, eta_min=1e-6). Prevents early divergence. |
| ✓Checkpoint Saving | Recovery | Saves best model (by val/fmax) to checkpoints/. Only the PaTz model state dict (not ESM-2) is persisted. |
| ✓nvidia-smi GPU Monitor | 0 VRAM overhead | API server uses subprocess nvidia-smi for GPU stats. No torch.cuda context held by the web server — all VRAM available for training. |
LR Schedule — Cosine Warm Restarts
T_0=10, T_mult=2, eta_min=1e-6. 5-epoch linear warmup.
Cosine Warm Restarts: warmup 0→5 epochs, then T_0=10 T_mult=2 cycles
Estimated VRAM Budget (batch=8, FP16)
Approx. breakdown. Total ~4.6 GB of 6 GB available.
ESM-2 Inference2.1 GB
Activations (FP16)1.4 GB
Optimizer State0.4 GB
Buffer/Misc0.3 GB
PaTz Weights0.2 GB
Gradients0.2 GB
Total: ~4.6 GB / 6 GB (77%) — 1.4 GB headroom
Optimizer Config (AdamW)
name:adamw
learning_rate:1.0e-04
betas:[0.9, 0.999]
eps:1.0e-08
weight_decay:1.0e-05
grad_clip:1.0
warmup_epochs:5
scheduler:cosine_with_restarts (T_0=10, T_mult=2)
early_stopping:monitor: val/functional_f1_macro, patience: 10