Tune VRAM memory: delete logits before backward, safe micro_batch 4 for 32GB GPUs

code/train_sft.py

@@ -169,13 +169,13 @@ def auto_profile_hardware():
     # Tier Classification for 2048 Context Length (Target Effective Batch = 128)
     if vram_gb >= 75:  # H200 (141GB), H100 (80GB), GH200, A100-80GB
         tier = "Ultra-Tier"
-        micro_batch = 16
-        grad_accum = 8
+        micro_batch = 8
+        grad_accum = 16
         desc = "NVIDIA Hopper / Datacenter Beast (141GB / 80GB HBM3e)"
     elif vram_gb >= 30:  # RTX 5090 (32GB), A100 (40GB), A6000 (48GB), RTX 6000 Ada
         tier = "High-Tier"
-        micro_batch = 8
-        grad_accum = 16
+        micro_batch = 4
+        grad_accum = 32
         desc = "NVIDIA Blackwell / High-End Workstation (32GB+)"
     elif vram_gb >= 20:  # RTX 4090 (24GB), RTX 3090 (24GB), L4 (24GB), A10G (24GB)
         tier = "Pro-Tier"
@@ -753,6 +753,7 @@ def main():
                         domain_loss_sum[d_id] += s_loss
                         domain_token_cnt[d_id] += act
 
+            del logits, shift_logits
             total_val_loss += sum_loss.item()
             val_tokens += num_active
             
@@ -802,6 +803,7 @@ def main():
                 logits, loss = model(inputs, targets=labels, pad_id=PAD_TOKEN_ID, shift_labels=True)
                 loss_scaled = loss / grad_acc
 
+            del logits
             loss_scaled.backward()
             accum_loss += loss.item()
             epoch_loss += loss.item()