Upload 2 files

main.py

@@ -92,14 +92,24 @@ print(dataset_data[0])
 with open("mitre-dataset.json", "w") as f:
     json.dump(dataset_data, f)
 
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+quantization_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True)
 
 BASE_MODEL = "decapoda-research/llama-7b-hf"
 
-model = LlamaForCausalLM.from_pretrained(
+device_map = {
+    "transformer.word_embeddings": 0,
+    "transformer.word_embeddings_layernorm": 0,
+    "lm_head": "cpu",
+    "transformer.h": 0,
+    "transformer.ln_f": 0,
+}
+
+model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    load_in_8bit=True,
-    torch_dtype=torch.float16,
     device_map="auto",
+    quantization_config=quantization_config,
 )
 
 tokenizer = LlamaTokenizer.from_pretrained(BASE_MODEL)
@@ -196,7 +206,6 @@ training_arguments = transformers.TrainingArguments(
     warmup_steps=100,
     max_steps=TRAIN_STEPS,
     learning_rate=LEARNING_RATE,
-    fp16=True,
     logging_steps=10,
     optim="adamw_torch",
     evaluation_strategy="steps",
@@ -228,7 +237,14 @@ model.state_dict = (
     )
 ).__get__(model, type(model))
 
+print("Compiling model...")
 model = torch.compile(model)
+print("Done compiling model...")
 
+print("Training model...")
 trainer.train()
+print("Done training model...")
+
+print("Saving model...")
 model.save_pretrained(OUTPUT_DIR)
+print("Done saving model...")

prepare.sh

@@ -3,6 +3,7 @@ curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.s
 apt-get install -y git
 apt-get install -y git-lfs
 apt-get install -y python3-pip
-git clone https://huggingface.co/khoicrtp/test_model
-cd test_model
-python3 finetune_lora.py
+git clone https://huggingface.co/khoicrtp/test_scratch
+cd test_scratch
+pip3 install -r requirements.txt --user
+python3 main.py