Ray fine-tuning example (#296)

gibchikafa · web-flow · commit 5cd30d12246e · 2025-01-10T14:06:01.000+01:00
diff --git a/llm-ai-systems/fine-tuning/configs/deepspeed_configs/zero_3_llama_2_13b.json b/llm-ai-systems/fine-tuning/configs/deepspeed_configs/zero_3_llama_2_13b.json
@@ -0,0 +1,35 @@
+{
+    "fp16": {
+        "enabled": "auto"
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": 5e8,
+        "stage3_prefetch_bucket_size": 5e8,
+        "stage3_param_persistence_threshold": 1e6,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": true,
+        "round_robin_gradients": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 10,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
diff --git a/llm-ai-systems/fine-tuning/configs/deepspeed_configs/zero_3_llama_2_70b.json b/llm-ai-systems/fine-tuning/configs/deepspeed_configs/zero_3_llama_2_70b.json
@@ -0,0 +1,28 @@
+{
+    "fp16": {
+        "enabled": false
+    },
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": false
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "gather_16bit_weights_on_model_save": true,
+        "round_robin_gradients": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 10,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
diff --git a/llm-ai-systems/fine-tuning/configs/deepspeed_configs/zero_3_llama_2_7b.json b/llm-ai-systems/fine-tuning/configs/deepspeed_configs/zero_3_llama_2_7b.json
@@ -0,0 +1,35 @@
+{   
+    "fp16": {
+        "enabled": "auto"
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": 5e8,
+        "stage3_prefetch_bucket_size": 5e8,
+        "stage3_param_persistence_threshold": 1e6,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": true,
+        "round_robin_gradients": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 10,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
diff --git a/llm-ai-systems/fine-tuning/configs/deepspeed_configs/zero_3_offload_optim_param.json b/llm-ai-systems/fine-tuning/configs/deepspeed_configs/zero_3_offload_optim_param.json
@@ -0,0 +1,32 @@
+{
+    "fp16": {
+        "enabled": "auto"
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": false
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": false
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "gather_16bit_weights_on_model_save": true,
+        "round_robin_gradients": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 10,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
diff --git a/llm-ai-systems/fine-tuning/configs/lora_configs/lora.json b/llm-ai-systems/fine-tuning/configs/lora_configs/lora.json
@@ -0,0 +1,11 @@
+{   
+    "r": 8,
+    "lora_alpha": 16,
+    "lora_dropout": 0.05,
+    "target_modules": ["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "embed_tokens", "lm_head"],
+    "task_type": "CAUSAL_LM",
+    "modules_to_save": [],
+    "bias": "none",
+    "fan_in_fan_out": false,
+    "init_lora_weights": true
+}
diff --git a/llm-ai-systems/fine-tuning/llama_fine_tune_runtime_env.yaml b/llm-ai-systems/fine-tuning/llama_fine_tune_runtime_env.yaml
@@ -0,0 +1,12 @@
+pip:
+  - transformers==4.44.0
+  - accelerate==0.31.0
+  - peft==0.11.1
+  - deepspeed==0.16.2
+env_vars:
+  LIBRARY_PATH: "$CUDA_HOME/lib64:$LIBRARY_PATH"
+  PROJECT_DIR: "/home/yarnapp/hopsfs"
+  TRAINED_MODEL_STORAGE_PATH: "${PROJECT_DIR}/Resources/llama_finetuning/fine-tuned-model"
+  TRAINING_DATA_DIR: "${PROJECT_DIR}/Resources/llama_finetuning/datasets"
+  TRAINING_CONFIGURATION_DIR: "${PROJECT_DIR}/Resources/llama_finetuning/configs"
+  
diff --git a/llm-ai-systems/fine-tuning/llama_fine_tuning_with_ray.ipynb b/llm-ai-systems/fine-tuning/llama_fine_tuning_with_ray.ipynb
diff --git a/llm-ai-systems/fine-tuning/ray_llm_finetuning.py b/llm-ai-systems/fine-tuning/ray_llm_finetuning.py
