e2e test worker

Author: earonesty

ai_worker/fine_tune.py

@@ -1,37 +1,59 @@
+import gc
 import base64
 import hashlib
 import json
+import asyncio
+import threading
 import logging
 import os
 import random
+import tarfile
+import shutil
 
 import transformers
 from datasets import load_dataset
 from httpx import AsyncClient, Response
 
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainerCallback 
 
-from ai_worker.jsonlines import load_jsonlines
+from gguf_loader.convert import main as gguf_main
 
 MAX_CONTEXT = 300000
 
 log = logging.getLogger(__name__)
 
+def gzip(folder):
+    """tar gz the folder to 'folder.tar.gz', removes the folder"""
+    base_folder_name = os.path.basename(folder)
+    with tarfile.open(f"{folder}.tar.gz", 'w:gz') as archive:
+        archive.add(folder, arcname=base_folder_name)
+    return f"{folder}.tar.gz"
 
 class FineTuner:
     def __init__(self, conf):
         self.conf = conf
         os.makedirs(self.conf.tmp_dir, exist_ok=True)
 
-    def temp_file(self, name):
-        return os.path.join(self.conf.tmp_dir, name)
+    def temp_file(self, name, wipe=False):
+        ret = os.path.join(self.conf.tmp_dir, name)
+        if wipe:
+            shutil.rmtree(ret, ignore_errors=True)
+        return ret
 
     def massage_line(self, ln, job):
+        # toss our role for now, for some reason it didn't work
+        # todo: check for role support in template
+        if "mistral" in job["model"].lower():
+            j = json.loads(ln)
+            j["messages"] = [m for m in j["messages"] if m["role"] != "system"]
+            ln = json.dumps(j) + "\n"
         return ln
 
     def massage_fine_tune(self, file, job):
         cnt = 0
+        tc = 0
+        ec = 0
         training_split_pct = job.get("hyperparameters", {}).get("training_split", 0.8)
 
         train_file = file + ".train"
@@ -41,42 +63,70 @@ def massage_fine_tune(self, file, job):
             with open(eval_file, "w") as ef:
                 with open(file, "r") as inp:
                     ln = inp.readline(MAX_CONTEXT)
-                    ln = self.massage_line(ln, job)
                     while ln:
+                        ln = self.massage_line(ln, job)
                         cnt += 1
-                        if random.random() > training_split_pct:
+                        if ec and (random.random() > training_split_pct or tc <= ec):
+                            tc += 1
                             tf.write(ln)
                         else:
+                            ec += 1
                             ef.write(ln)
                         ln = inp.readline(MAX_CONTEXT)
-                        ln = self.massage_line(ln, job)
         return train_file, eval_file
 
     async def fine_tune(self, job):
         log.info("fine tuning: %s", job)
 
-        yield {"status": "downloading_data"}
+        yield {"status": "download_data"}
 
-        base_model = job["model"]
         training_url = job["training_file"]
         training_file = await self.download_file(training_url)
-
+        job["training_file"] = training_file
+
+        q = asyncio.Queue()
+       
+        loop = asyncio.get_running_loop()
+
+        t = threading.Thread(target=lambda: self._fine_tune(job, lambda res: loop.call_soon_threadsafe(q.put_nowait, res)), daemon=True)
+        
+        t.start()
+        while True:
+            res = await q.get()
+            if res is None:
+                break
+            yield res
+        log.info("DONE")
+        t.join()
+
+    def _fine_tune(self, job, cb):
+        try:
+            self._unsafe_fine_tune(job, cb)
+        except Exception as ex:
+            log.exception("error in fine tune")
+            cb({"status": "error", "detail": repr(ex)})
+        finally:
+            cb(None)
+
+    def _unsafe_fine_tune(self, job, cb):
+        training_file = job["training_file"]
         train_file, eval_file = self.massage_fine_tune(training_file, job)
 
-        train_dataset = load_jsonlines(open(train_file))
-        eval_dataset = load_jsonlines(open(eval_file))
-
-        # todo: use user's model request
+        base_model = job["model"]
+        datasets = load_dataset("json", data_files={"train": train_file, "eval": eval_file})
+        train_dataset = datasets["train"]
+        eval_dataset = datasets["eval"]
 
-        base_model_id = "mistralai/Mistral-7B-v0.1"
+        base_model_id = base_model.split(":")[0]
 
         # todo: use hyperparams and Q_ filter, if present, for this
 
         hp = job.get("hyperparameters", {})
 
         args = {}
 
-        yield {"status": "loading_model"}
+        log.info("load model")
+        cb({"status": "load_model"})
 
         args.update(dict(
             load_in_4bit=True,
@@ -89,25 +139,23 @@ async def fine_tune(self, job):
 
         # todo: ideally we use llama cpp, but the cuda support for finetune isn't there
 
-        model = AutoModelForCausalLM.from_pretrained(base_model_id, quantization_config=bnb_config, device_map="auto")
-
         tokenizer = AutoTokenizer.from_pretrained(
             base_model_id,
             padding_side="left",
             add_eos_token=True,
             add_bos_token=True,
         )
 
-        train_dataset = tokenizer.apply_chat_template(train_dataset)
-        eval_dataset = tokenizer.apply_chat_template(eval_dataset)
-
+        # sadly, does not take generators, just loads everything in ram
         tokenizer.pad_token = tokenizer.eos_token
-
-        max_length = 512
-
+        # todo: derive from model params
+        max_length = 4096
         def generate_and_tokenize_prompt(prompt):
+            # all input is openai formatted, and we clean it up above if needed
+            pr = prompt["messages"]
+            tmpl = tokenizer.apply_chat_template(pr, tokenize=False)
             result = tokenizer(
-                prompt,
+                tmpl,
                 truncation=True,
                 max_length=max_length,
                 padding="max_length",
@@ -118,7 +166,9 @@ def generate_and_tokenize_prompt(prompt):
         tokenized_train_dataset = train_dataset.map(generate_and_tokenize_prompt)
         tokenized_val_dataset = eval_dataset.map(generate_and_tokenize_prompt)
 
-        from peft import prepare_model_for_kbit_training
+        model = AutoModelForCausalLM.from_pretrained(base_model_id, quantization_config=bnb_config, device_map="auto", resume_download=True)
+
+        from peft import prepare_model_for_kbit_training, PeftModel
 
         model.gradient_checkpointing_enable()
         model = prepare_model_for_kbit_training(model)
@@ -127,7 +177,7 @@ def generate_and_tokenize_prompt(prompt):
 
         config = LoraConfig(
             r=32,
-            lora_alpha=64,
+            lora_alpha=hp.get("lora_alpha", 64),
             target_modules=[
                 "q_proj",
                 "k_proj",
@@ -139,7 +189,7 @@ def generate_and_tokenize_prompt(prompt):
                 "lm_head",
             ],
             bias="none",
-            lora_dropout=0.05,  # Conventional
+            lora_dropout=hp.get("lora_dropout", 0.05),  # Conventional
             task_type="CAUSAL_LM",
         )
 
@@ -162,41 +212,114 @@ def generate_and_tokenize_prompt(prompt):
             model.model_parallel = True
 
         project = "journal-finetune"
-        base_model_name = "mistral"
-        run_name = base_model_name + "-" + project
+        base_model_name = base_model_id.split("/")[-1]
+        run_name = base_model_name + "-" + project + "-" + os.urandom(16).hex()
         output_dir = "./" + run_name
 
         tokenizer.pad_token = tokenizer.eos_token
 
+        class EarlyStoppingCallback(TrainerCallback):
+            def on_log(self, args, state, control, logs=None, **kwargs):
+                cb({"status": "log", "logs": logs})
+                eval_loss = logs.get("eval_loss", None)
+                if eval_loss is not None and eval_loss <= hp.get("stop_eval_loss", 0.05):
+                    print("Early stopping criterion reached!")
+                    control.should_training_stop = True
+
+            def on_save(self, args, state, control, **kwargs):
+                checkpoint_dir = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
+                log.info(f"checkpoint {checkpoint_dir}")
+                cb({"status": "checkpoint"})
+
+
         trainer = transformers.Trainer(
             model=model,
             train_dataset=tokenized_train_dataset,
             eval_dataset=tokenized_val_dataset,
+            callbacks=[EarlyStoppingCallback()],
             args=transformers.TrainingArguments(
                 output_dir=output_dir,
                 warmup_steps=1,
-                per_device_train_batch_size=2,
-                gradient_accumulation_steps=1,
-                max_steps=500,
-                learning_rate=2.5e-5,  # Want a small lr for finetuning
+                per_device_train_batch_size=hp.get("batch_size", 4),
+                gradient_accumulation_steps=hp.get("accumulation_steps", 4),
+                max_steps=hp.get("max_steps", -1),
+                num_train_epochs=hp.get("n_epochs", 3), # use openai terminology here
+                learning_rate=hp.get("learning_rate_multiplier", 2.5e-5),  # Want a small lr for finetuning
                 bf16=True,
                 optim="paged_adamw_8bit",
                 logging_steps=25,  # When to start reporting loss
                 logging_dir="./logs",  # Directory for storing logs
                 save_strategy="steps",  # Save the model checkpoint every logging step
                 save_steps=25,  # Save checkpoints
+                save_total_limit=5,  # Save checkpoints
+                load_best_model_at_end=True,
                 evaluation_strategy="steps",  # Evaluate the model every logging step
-                eval_steps=25,  # Evaluate and save checkpoints every 50 steps
+                eval_steps=25,  # Evaluate and save checkpoints every 25 steps
                 do_eval=True,  # Perform evaluation at the end of training
             ),
             data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
         )
 
-        model.config.use_cache = False  # silence the warnings. Please re-enable for inference!
+        log.info("start train")
+        cb({"status": "start_train"})
+        model.config.use_cache = False  # silence the warnings
+        
         trainer.train()
 
-        res = {"status": "done", "checkpoint": str(base64.b64encode(b"checkpoint"))}
-        yield res
+        tmp = self.temp_file(run_name, wipe=True)
+        tokenizer.save_pretrained(tmp)
+
+        self.return_final(run_name, model, cb)
+
+    def return_final(self, run_name, model, cb):
+        log.info("return final")
+
+        tmp = self.temp_file(run_name)
+        
+        # send up lora
+        model.save_pretrained(tmp, safe_serialization=True)
+        gz = gzip(tmp)
+        shutil.rmtree(tmp)
+        with open(gz, "rb") as fil:
+            while True:
+                dat = fil.read(100000)
+                if not dat:
+                    break;
+                res = {"status": "lora", "chunk": str(base64.b64encode(dat))}
+                cb(res)
+      
+        log.info("merge weights")
+
+        # merge weights
+        
+        # reload as float16 for merge
+        del model
+        gc.collect()
+
+        model = PeftModel.from_pretrained(AutoModelForCausalLM.from_pretrained(base_model_id, torch_dtype=torch.float16, local_files_only=True, device_map="auto"), tmp)
+        model = model.merge_and_unload()
+        
+        gc.collect()
+        model.save_pretrained(tmp)
+        
+        # convert to gguf for fast inference
+        log.info("ggml convert")
+        gguf_main([tmp])
+        gg = tmp + "/ggml-model-f16.gguf"
+        with open(gg, "rb") as fil:
+            while True:
+                dat = fil.read(100000)
+                if not dat:
+                    break;
+                res = {"status": "gguf", "chunk": str(base64.b64encode(dat))}
+                cb(res)
+        
+        shutil.rmtree(tmp)
+        shutil.rmtree(output_dir)
+        
+        res = {"status": "done"}
+        log.info("done train")
+        cb(res)
 
     async def download_file(self, training_url: str) -> str:
         output_file = self.temp_file(hashlib.md5(training_url.encode()).hexdigest())

ai_worker/jsonlines.py

@@ -4,5 +4,7 @@
 def load_jsonlines(fin):
     while True:
         lin = fin.readline()
+        if not lin:
+            return
         yield json.loads(lin)
 

build-bin.sh

@@ -9,6 +9,12 @@ if [ -z "$cmake" -o -z "$gpu" ]; then
     exit 1
 fi
 
+with_torch=""
+if [ "$gpu" == "cuda-torch" ]; then
+    with_torch="--with torch"
+fi
+
+
 set -o xtrace
 
 python -mvenv "build-$gpu"
@@ -22,7 +28,7 @@ pip uninstall -y llama-cpp-python
 rm -f ~/AppData/Local/pypoetry/Cache/artifacts/*/*/*/*/llama*
 rm -f ~/.cache/pypoetry/artifacts/*/*/*/*/llama*
 
-CMAKE_ARGS="$cmake" FORCE_CMAKE=1 poetry install
+CMAKE_ARGS="$cmake" FORCE_CMAKE=1 poetry install $with_torch
 
 python build-version.py
 

build-linux.sh

@@ -6,4 +6,6 @@ set -o xtrace
 
 ./build-bin.sh opencl linux-64 "-DLLAMA_CLBLAST=ON" 
 
+./build-bin.sh cuda-torch linux-64 "-DLLAMA_CLBLAST=ON" 
+
 ./upload.sh