Optimize memory usage for T4 GPU training

- Reduce max_seq_length from 2048 to 1024
- Reduce LoRA rank from 16 to 8 (fewer trainable params)
- Enable gradient checkpointing on model
- Add gradient_checkpointing=True in TrainingArguments
- Reduce gradient_accumulation_steps from 16 to 8
- Use paged_adamw_8bit optimizer for GPU memory efficiency

These changes should prevent OOM crashes on T4 (16GB VRAM)

finetune.py

@@ -13,7 +13,7 @@ def main():
     # Configuration
     model_name = "Qwen/Qwen2.5-0.5B-Instruct"  # Using 0.5B as 0.6B doesn't exist
     output_dir = "./qwen-codeforces-cots"
-    max_seq_length = 2048
+    max_seq_length = 1024  # Reduced from 2048 to save memory
 
     # Detect device - prefer CUDA for GPU training
     if torch.cuda.is_available():
@@ -63,17 +63,20 @@ def main():
         )
         from peft import prepare_model_for_kbit_training
         model = prepare_model_for_kbit_training(model)
+        # Enable gradient checkpointing for memory efficiency
+        model.gradient_checkpointing_enable()
     else:
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
             torch_dtype=torch.float32,
             trust_remote_code=True,
         )
+        model.gradient_checkpointing_enable()
 
-    # LoRA config
+    # LoRA config - reduced rank for memory efficiency
     lora_config = LoraConfig(
-        r=16,
-        lora_alpha=32,
+        r=8,  # Reduced from 16 to save memory
+        lora_alpha=16,  # Reduced proportionally
         target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
         lora_dropout=0.05,
         bias="none",
@@ -122,16 +125,17 @@ def main():
         mlm=False,  # We're doing causal LM, not masked LM
     )
 
-    # Training arguments - reduced for CPU training
+    # Training arguments - optimized for T4 GPU
     training_args = TrainingArguments(
         output_dir=output_dir,
-        per_device_train_batch_size=1,  # Reduced for CPU
+        per_device_train_batch_size=1,  # Keep at 1 for memory safety
         per_device_eval_batch_size=1,
-        gradient_accumulation_steps=16,  # Maintain effective batch size
+        gradient_accumulation_steps=8,  # Reduced from 16 to lower memory pressure
         num_train_epochs=1,
         max_steps=1000,  # Limit steps for testing
         learning_rate=2e-4,
         fp16=use_fp16,
+        gradient_checkpointing=True,  # Enable gradient checkpointing to save memory
         save_strategy="steps",
         save_steps=200,  # Save more frequently
         eval_strategy="steps",
@@ -139,7 +143,7 @@ def main():
         logging_steps=10,
         warmup_steps=50,
         lr_scheduler_type="cosine",
-        optim="adamw_torch",
+        optim="paged_adamw_8bit" if torch.cuda.is_available() else "adamw_torch",  # Use 8-bit optimizer on GPU
         report_to="none",
         max_grad_norm=0.3,
         save_total_limit=2,