Upload qlora.py

code/qlora.py

@@ -0,0 +1,183 @@
+import torch, os, wandb, uuid, json
+import bitsandbytes as bnb
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, BitsAndBytesConfig, TrainerCallback
+from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from datasets import load_dataset, DatasetDict, Dataset,load_from_disk
+from functools import partial
+
+set_seed(42)
+
+accelerator = Accelerator()
+run_id = str(uuid.uuid4())
+modelpath="microsoft/phi-2"
+dataset_name="teknium/OpenHermes-2.5"
+lr=0.00002
+bs=10            # batch size
+bs_eval=16      # batch size for evals
+ga_steps=4     # gradient acc. steps
+epochs=1
+max_length=1024
+output_dir=f"out_{run_id}"
+
+# Load model
+model = AutoModelForCausalLM.from_pretrained(
+    modelpath,    
+    device_map={"": accelerator.process_index},
+    quantization_config=BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_quant_type="nf4",
+    ),
+    torch_dtype=torch.bfloat16,
+    # does not work yet
+    # attn_implementation="flash_attention_2",          
+)
+
+# Load Tokenizer
+tokenizer = AutoTokenizer.from_pretrained(modelpath, use_fast=False)    # fast tokenizer sometimes ignores the added tokens
+
+# Add tokens <|im_start|> and <|im_end|>, latter is special eos token, 
+tokenizer.add_tokens(["<|im_start|>", "<PAD>"])
+tokenizer.pad_token = "<PAD>"
+tokenizer.add_special_tokens(dict(eos_token="<|im_end|>"))
+model.config.eos_token_id = tokenizer.eos_token_id
+
+# Add adapters to model
+model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True) 
+
+lora_config = LoraConfig(
+    r=32, 
+    lora_alpha=32, 
+    target_modules = [ "q_proj", "k_proj", "v_proj", "dense" ],
+    modules_to_save = ["lm_head", "embed_tokens"],
+    lora_dropout=0.1, 
+    bias="none", 
+    task_type="CAUSAL_LM",
+)
+model = get_peft_model(model, lora_config)
+
+model.config.use_cache = False
+
+# Print stats
+if accelerator.is_main_process:
+    model.print_trainable_parameters()
+
+# Load dataset
+with accelerator.main_process_first():
+    dataset = load_dataset(dataset_name)
+    dataset = dataset["train"].train_test_split(test_size=0.1)
+
+    # Format (chatML) and tokenize dataset
+    templates= {
+        "system": "<|im_start|>system\n{msg}<|im_end|>",
+        "human": "<|im_start|>user\n{msg}<|im_end|>",
+        "gpt": "<|im_start|>assistant\n{msg}<|im_end|>",
+    }
+    IGNORE_INDEX=-100
+
+    def tokenize(input, max_length):
+        input_ids, attention_mask, labels = [], [], []
+
+        for i,msg in enumerate(input["conversations"]):
+            msg_role=msg["from"]
+            msg_content=msg["value"]
+            isHuman=msg_role=="human"
+            if not msg_role in templates: return  # this will break it
+            msg_chatml=templates[msg_role].format(msg=msg_content)
+            msg_tokenized=tokenizer(msg_chatml, truncation=False, add_special_tokens=False)
+
+            input_ids+=msg_tokenized["input_ids"]
+            attention_mask+=msg_tokenized["attention_mask"]
+            labels+=[IGNORE_INDEX]*len(msg_tokenized["input_ids"]) if isHuman else msg_tokenized["input_ids"]
+
+        return {
+            "input_ids": input_ids[:max_length],
+            "attention_mask": attention_mask[:max_length],
+            "labels": labels[:max_length],
+        }
+
+    dataset_tokenized = dataset.map(
+        partial(tokenize, max_length=max_length), 
+        batched=False, 
+        # num_proc=os.cpu_count()//accelerator.num_processes,    # multithreaded
+        num_proc=os.cpu_count(),    # multithreaded
+        remove_columns=dataset["train"].column_names  # don't need this anymore, we have tokens from here on
+    )
+
+# collate function - to transform list of dictionaries [ {input_ids: [123, ..]}, {.. ] to single batch dictionary { input_ids: [..], labels: [..], attention_mask: [..] }
+def collate(elements):
+    tokens=[e["input_ids"] for e in elements]
+    tokens_maxlen=max([len(t) for t in tokens])
+
+    for i,sample in enumerate(elements):
+        input_ids=sample["input_ids"]
+        labels=sample["labels"]
+        attention_mask=sample["attention_mask"]
+
+        pad_len=tokens_maxlen-len(input_ids)
+
+        input_ids.extend( pad_len * [tokenizer.pad_token_id] )   
+        labels.extend( pad_len * [IGNORE_INDEX] )    
+        attention_mask.extend( pad_len * [0] ) 
+
+    batch={
+        "input_ids": torch.tensor( [e["input_ids"] for e in elements] ),
+        "labels": torch.tensor( [e["labels"] for e in elements] ),
+        "attention_mask": torch.tensor( [e["attention_mask"] for e in elements] ),
+    }
+
+    return batch
+ 
+steps_per_epoch=len(dataset_tokenized["train"])//(accelerator.num_processes*bs*ga_steps)
+
+args = TrainingArguments(
+    output_dir=output_dir,
+    per_device_train_batch_size=bs,
+    per_device_eval_batch_size=bs_eval,
+    evaluation_strategy="steps",
+    logging_steps=1,
+    eval_steps=steps_per_epoch//3,    # 2 evals per epoch
+    save_steps=steps_per_epoch//3,     # save once per epoch
+    gradient_accumulation_steps=ga_steps,
+    num_train_epochs=epochs,
+    lr_scheduler_type="constant",
+    optim="paged_adamw_32bit",      # val_loss will go nan with paged_adamw_8bit
+    learning_rate=lr,
+    group_by_length=False,
+    bf16=True,        
+    ddp_find_unused_parameters=False,
+)
+
+trainer = Trainer(
+    model=model,
+    tokenizer=tokenizer,
+    args=args,
+    data_collator=collate,
+    train_dataset=dataset_tokenized["train"],
+    eval_dataset=dataset_tokenized["test"],
+)
+
+if accelerator.is_main_process:
+    run = wandb.init(
+        project="phi2-teknium1",
+        name=modelpath+"_"+dataset_name+f"_bs-{bs}_LR-{lr}_GPUs-{accelerator.num_processes}_maxlen-{max_length}_{run_id}",
+        config={
+            "model_name": modelpath,
+            "run_id": run_id,
+            "dataset": dataset_name,
+            "output_dir": output_dir,
+            "lr": lr,
+            "max_length": max_length,
+            "train_batch_size": bs,
+            "validation_batch_size": bs,
+            "ga_steps": ga_steps,
+            "lora_config": lora_config, 
+            "training_args": args,
+            "GPUs": accelerator.num_processes,
+        }
+    )
+    run.log_code()
+
+trainer.train()