Automate LLM Finetune Deepspeed demo

Author: abhijeet-dhumal

go.mod

@@ -16,6 +16,8 @@ require (
 	sigs.k8s.io/kueue v0.6.2
 )
 
+replace github.com/project-codeflare/codeflare-common => /home/abdhumal/abhidev/RedHatDev/codeflare-common
+
 require (
 	github.com/aymerick/douceur v0.2.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect

tests/odh/ray_finetune_llm_deepspeed_test.go

@@ -0,0 +1,102 @@
+/*
+Copyright 2023.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package odh
+
+import (
+	"testing"
+
+	. "github.com/onsi/gomega"
+	. "github.com/project-codeflare/codeflare-common/support"
+	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+)
+
+func TestRayFinetuneDemo(t *testing.T) {
+	mnistRayLlmFinetune(t, 1)
+}
+
+func mnistRayLlmFinetune(t *testing.T, numGpus int) {
+	test := With(t)
+
+	// Create a namespace
+	namespace := test.NewTestNamespace()
+
+	// Test configuration
+	jupyterNotebookConfigMapFileName := "ray_finetune_llm_deepspeed.ipynb"
+
+	// Test configuration
+	configMap := map[string][]byte{
+		// MNIST Ray Notebook
+		jupyterNotebookConfigMapFileName: ReadFile(test, "resources/ray_finetune_demo/ray_finetune_llm_deepspeed.ipynb"),
+		"ray_finetune_llm_deepspeed.py":  ReadFile(test, "resources/ray_finetune_demo/ray_finetune_llm_deepspeed.py"),
+		"ray_finetune_requirements.txt":  ReadRayFinetuneRequirementsTxt(test),
+		"create_dataset.py":              ReadFile(test, "resources/ray_finetune_demo/create_dataset.py"),
+		"lora.json":                      ReadFile(test, "resources/ray_finetune_demo/lora.json"),
+		"zero_3_llama_2_7b.json":         ReadFile(test, "resources/ray_finetune_demo/zero_3_llama_2_7b.json"),
+		"utils.py":                       ReadFile(test, "resources/ray_finetune_demo/utils.py"),
+	}
+
+	config := CreateConfigMap(test, namespace.Name, configMap)
+
+	// Define the regular(non-admin) user
+	userName := GetNotebookUserName(test)
+	userToken := GetNotebookUserToken(test)
+
+	// Create role binding with Namespace specific admin cluster role
+	CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin")
+
+	// Create Notebook CR
+	createNotebook(test, namespace, userToken, config.Name, jupyterNotebookConfigMapFileName, numGpus)
+
+	// Gracefully cleanup Notebook
+	defer func() {
+		deleteNotebook(test, namespace)
+		test.Eventually(listNotebooks(test, namespace), TestTimeoutGpuProvisioning).Should(HaveLen(0))
+	}()
+
+	// Make sure the RayCluster is created and running
+	test.Eventually(RayClusters(test, namespace.Name), TestTimeoutGpuProvisioning).
+		Should(
+			And(
+				HaveLen(1),
+				ContainElement(WithTransform(RayClusterState, Equal(rayv1.Ready))),
+			),
+		)
+
+	// Make sure the RayCluster finishes and is deleted
+	test.Eventually(RayClusters(test, namespace.Name), TestTimeoutGpuProvisioning).
+		Should(HaveLen(0))
+}
+
+func ReadRayFinetuneRequirementsTxt(test Test) []byte {
+	// Read the requirements.txt from resources and perform replacements for custom values using go template
+	props := struct {
+		PipIndexUrl    string
+		PipTrustedHost string
+	}{
+		PipIndexUrl: "--index " + string(GetPipIndexURL()),
+	}
+
+	// Provide trusted host only if defined
+	if len(GetPipTrustedHost()) > 0 {
+		props.PipTrustedHost = "--trusted-host " + GetPipTrustedHost()
+	}
+
+	template, err := files.ReadFile("resources/ray_finetune_demo/ray_finetune_requirements.txt")
+	test.Expect(err).NotTo(HaveOccurred())
+
+	return ParseTemplate(test, template, props)
+}

tests/odh/resources/ray_finetune_demo/create_dataset.py

@@ -0,0 +1,31 @@
+from datasets import load_dataset
+import json
+import os
+
+dataset = load_dataset("gsm8k", "main")
+
+dataset_splits = {"train": dataset["train"], "test": dataset["test"]}
+
+
+def main():
+    if not os.path.exists("data"):
+        os.mkdir("data")
+
+    with open("data/tokens.json", "w") as f:
+        tokens = {}
+        tokens["tokens"] = ["<START_Q>", "<END_Q>", "<START_A>", "<END_A>"]
+        f.write(json.dumps(tokens))
+
+    for key, ds in dataset_splits.items():
+        with open(f"data/{key}.jsonl", "w") as f:
+            for item in ds:
+                newitem = {}
+                newitem["input"] = (
+                    f"<START_Q>{item['question']}<END_Q>"
+                    f"<START_A>{item['answer']}<END_A>"
+                )
+                f.write(json.dumps(newitem) + "\n")
+
+
+if __name__ == "__main__":
+    main()

tests/odh/resources/ray_finetune_demo/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
+}