feat: use DeepSpeed ZeRO-3 for optimal 2-GPU training

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

@@ -46,12 +46,22 @@ 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
+# Distributed training (2x RTX 5090) - DeepSpeed ZeRO-3
+plugin: deepspeed
+huggingface_hub:
+  token: null
+deeepspeed_config:
+  zero_optimization:
+    stage: 3
+    offload_optimizer:
+      device: cpu
+      pin_memory: true
+    offload_param:
+      device: cpu
+      pin_memory: true
+  gradient_clipping: 1.0
+  train_batch_size: auto
+  train_micro_batch_size_per_gpu: auto
 
 # Evaluation
 eval_strategy: steps