Reduce context 4096->3072, add expandable_segments for OOM fix

train_dpo.py

@@ -22,6 +22,8 @@ import sys
 import torch
 from typing import Optional
 
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
 # ============================================================
 # CONFIG
 # ============================================================
@@ -57,8 +59,8 @@ class Config:
     per_device_train_batch_size = 1
     gradient_accumulation_steps = 4
     num_train_epochs = 1
-    max_length = 4096
-    max_prompt_length = 3072
+    max_length = 3072
+    max_prompt_length = 2304
     logging_steps = 10
     save_steps = 200
     eval_steps = 200