Add KFTO GPU tests for llama2 and granite using LoRA

opendatahub-io · Jul 16, 2024 · 296c1d6 · 296c1d6
1 parent 6e098bf
commit 296c1d6
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 6 deletions.
diff --git a/tests/kfto/core/config_granite_34b_code_instruct_lora.json b/tests/kfto/core/config_granite_34b_code_instruct_lora.json
@@ -0,0 +1,21 @@
+{
+    "model_name_or_path": "ibm-granite/granite-34b-code-instruct",
+    "training_data_path": "/mnt/scratch/dataset/alpaca_data.json",
+    "output_dir": "/mnt/output/model",
+    "num_train_epochs": 1.0,
+    "per_device_train_batch_size": 1,
+    "per_device_eval_batch_size": 4,
+    "gradient_accumulation_steps": 4,
+    "evaluation_strategy": "no",
+    "save_strategy": "epoch",
+    "learning_rate": 1e-5,
+    "weight_decay": 0.0,
+    "lr_scheduler_type": "cosine",
+    "include_tokens_per_second": true,
+    "response_template": "\n### Response:",
+    "dataset_text_field": "output",
+    "use_flash_attn": false,
+    "tokenizer_name_or_path": "ibm-granite/granite-34b-code-instruct",
+    "peft_method": "lora",
+    "target_modules": ["all-linear"]
+}
diff --git a/tests/kfto/core/config_llama2_13b.json → .../kfto/core/config_llama2_13b_chat_hf.json b/tests/kfto/core/config_llama2_13b.json → .../kfto/core/config_llama2_13b_chat_hf.json
diff --git a/tests/kfto/core/config_llama2_13b_chat_hf_lora.json b/tests/kfto/core/config_llama2_13b_chat_hf_lora.json
@@ -0,0 +1,20 @@
+{
+    "model_name_or_path": "meta-llama/Llama-2-13b-chat-hf",
+    "training_data_path": "/mnt/scratch/dataset/alpaca_data.json",
+    "output_dir": "/mnt/output/model",
+    "num_train_epochs": 1.0,
+    "per_device_train_batch_size": 1,
+    "per_device_eval_batch_size": 4,
+    "gradient_accumulation_steps": 4,
+    "evaluation_strategy": "no",
+    "save_strategy": "epoch",
+    "learning_rate": 1e-5,
+    "weight_decay": 0.0,
+    "lr_scheduler_type": "cosine",
+    "include_tokens_per_second": true,
+    "response_template": "\n### Response:",
+    "dataset_text_field": "output",
+    "use_flash_attn": false,
+    "tokenizer_name_or_path": "meta-llama/Llama-2-13b-chat-hf",
+    "peft_method": "lora"
+}
diff --git a/tests/kfto/core/kfto_kueue_sft_GPU_test.go b/tests/kfto/core/kfto_kueue_sft_GPU_test.go
@@ -39,22 +39,30 @@ func TestMultiGpuPytorchjobGranite20bCodeInstruct(t *testing.T) {
 	runMultiGpuPytorchjob(t, "config_granite_20b_code_instruct.json")
 }
 
-func TestMultiGpuPytorchjobLlama213b(t *testing.T) {
-	runMultiGpuPytorchjob(t, "config_llama2_13b.json")
+func TestMultiGpuPytorchjobGranite34bCodeInstructLoRa(t *testing.T) {
+	runMultiGpuPytorchjob(t, "config_granite_34b_code_instruct_lora.json")
 }
 
-func TestMultiGpuPytorchjobMetaLlama38bInstruct(t *testing.T) {
-	runMultiGpuPytorchjob(t, "config_meta_llama3_8b_instruct.json")
+func TestMultiGpuPytorchjobLlama213bChatHf(t *testing.T) {
+	runMultiGpuPytorchjob(t, "config_llama2_13b_chat_hf.json")
 }
 
-func TestMultiGpuPytorchjobMixtral8x7bInstructv01(t *testing.T) {
-	runMultiGpuPytorchjob(t, "config_mixtral_8x7b_instruct_v01.json")
+func TestMultiGpuPytorchjobLlama213bChatHfLoRa(t *testing.T) {
+	runMultiGpuPytorchjob(t, "config_llama2_13b_chat_hf_lora.json")
+}
+
+func TestMultiGpuPytorchjobMetaLlama38bInstruct(t *testing.T) {
+	runMultiGpuPytorchjob(t, "config_meta_llama3_8b_instruct.json")
 }
 
 func TestMultiGpuPytorchjobMetaLlama370bInstructLoRa(t *testing.T) {
 	runMultiGpuPytorchjob(t, "config_meta_llama3_70b_instruct_lora.json")
 }
 
+func TestMultiGpuPytorchjobMixtral8x7bInstructv01(t *testing.T) {
+	runMultiGpuPytorchjob(t, "config_mixtral_8x7b_instruct_v01.json")
+}
+
 func runMultiGpuPytorchjob(t *testing.T, modelConfigFile string) {
 	test := With(t)