fix: use FSDP for 2-GPU training, reduce batch size

training/configs/llama2-7b-lora.yaml

@@ -31,8 +31,8 @@ dataset:
 # Training Parameters
 train_params:
   num_train_epochs: 3
-  per_device_train_batch_size: 4
-  gradient_accumulation_steps: 4
+  per_device_train_batch_size: 1
+  gradient_accumulation_steps: 8
   learning_rate: 2e-4
   lr_scheduler_type: cosine
   weight_decay: 0.01
@@ -46,6 +46,13 @@ train_params:
 # Precision
 mixed_precision: bf16
 
+# Distributed training (2x RTX 5090)
+fsdp: full_shard
+fsdp_config:
+  limit_all_gathers: true
+  offload_optimizer: true
+  offload_model: false
+
 # Evaluation
 eval_strategy: steps
 eval_steps: 100