Upload sft_lora.py

supervised finetuninig by lora

sft_lora.py

@@ -0,0 +1,167 @@
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+from datasets import load_dataset
+import transformers
+from transformers import Trainer, TrainingArguments
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import BitsAndBytesConfig
+from peft import (
+    LoraConfig,
+    get_peft_model,
+    prepare_model_for_kbit_training,
+    set_peft_model_state_dict,
+)
+import torch
+
+
+
+CUTOFF_LEN = 1024  # 1024 accounts for about 96% of the data
+VAL_SET_SIZE = 2000
+DATA_PATH = "./dataset/Belle_open_source_0.5M.json" ## Choose dataset
+OUTPUT_DIR = "baichuansft"
+resume_from_checkpoint = "baichuansft"
+
+
+device_map = {"": 0}
+tokenizer = AutoTokenizer.from_pretrained("./baichuan-7B",trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("./baichuan-7B",
+                                             trust_remote_code=True,
+                                             quantization_config=BitsAndBytesConfig(
+                                                 load_in_4bit=True,
+                                                 bnb_4bit_compute_dtype=torch.bfloat16,
+                                                 bnb_4bit_use_double_quant=True,
+                                                 bnb_4bit_quant_type='nf4'
+                                             ),
+                                             device_map=device_map)
+
+model = prepare_model_for_kbit_training(model)
+
+### 所有的线性layer都装配上lora
+import bitsandbytes as bnb
+def find_all_linear_names(model):
+    #cls = bnb.nn.Linear8bitLt 
+    cls = bnb.nn.Linear4bit 
+    lora_module_names = set()
+    for name, module in model.named_modules():
+        if isinstance(module, cls):
+            names = name.split('.')
+            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
+
+
+    if 'lm_head' in lora_module_names: # needed for 16-bit
+        lora_module_names.remove('lm_head')
+    return list(lora_module_names)
+modules = find_all_linear_names(model)
+
+
+config = LoraConfig(
+    r=8,
+    lora_alpha=16,
+    lora_dropout=0.05,
+    bias="none",
+    target_modules=modules,
+    task_type="CAUSAL_LM",
+)
+
+
+model = get_peft_model(model, config)
+tokenizer.pad_token_id = 0
+
+if resume_from_checkpoint:
+    # Check the available weights and load them
+    checkpoint_name = os.path.join(
+        resume_from_checkpoint, "pytorch_model.bin"
+    )  # Full checkpoint
+    if not os.path.exists(checkpoint_name):
+        checkpoint_name = os.path.join(
+            resume_from_checkpoint, "adapter_model.bin"
+        )  # only LoRA model - LoRA config above has to fit
+        resume_from_checkpoint = (
+            False  # So the trainer won't try loading its state
+        )
+    # The two files above have a different name depending on how they were saved, but are actually the same.
+    if os.path.exists(checkpoint_name):
+        print(f"Restarting from {checkpoint_name}")
+        adapters_weights = torch.load(checkpoint_name)
+        set_peft_model_state_dict(model, adapters_weights)
+    else:
+        print(f"Checkpoint {checkpoint_name} not found")
+
+
+data = load_dataset("json", data_files=DATA_PATH)
+
+def tokenize(prompt, add_eos_token=True):
+    result = tokenizer(
+        prompt,
+        truncation=True,
+        max_length=CUTOFF_LEN,
+        padding=False,
+        return_tensors=None,
+    )
+    if (
+            result["input_ids"][-1] != tokenizer.eos_token_id
+            and len(result["input_ids"]) < CUTOFF_LEN
+            and add_eos_token
+    ):
+        result["input_ids"].append(tokenizer.eos_token_id)
+        result["attention_mask"].append(1)
+
+    if add_eos_token and len(result["input_ids"]) >= CUTOFF_LEN:
+        result["input_ids"][CUTOFF_LEN - 1] = tokenizer.eos_token_id
+        result["attention_mask"][CUTOFF_LEN - 1] = 1
+
+    result["labels"] = result["input_ids"].copy()
+
+    return result
+
+
+def generate_and_tokenize_prompt(data_point):
+    instruction = data_point['instruction']
+    input_text = data_point["input"]
+    input_text = "Human: " + instruction + input_text + "\n\nAssistant: "
+    input_text = tokenizer.bos_token + input_text if tokenizer.bos_token != None else input_text
+    target_text = data_point["output"] + tokenizer.eos_token
+    full_prompt = input_text + target_text
+    tokenized_full_prompt = tokenize(full_prompt)
+    return tokenized_full_prompt
+
+
+if VAL_SET_SIZE > 0:
+    train_val = data["train"].train_test_split(
+        test_size=VAL_SET_SIZE, shuffle=True, seed=42
+    )
+    train_data = train_val["train"].shuffle().map(generate_and_tokenize_prompt)
+    val_data = train_val["test"].shuffle().map(generate_and_tokenize_prompt)
+else:
+    train_data = data['train'].shuffle().map(generate_and_tokenize_prompt)
+    val_data = None
+
+trainer = Trainer(
+    model=model,
+    train_dataset=train_data,
+    eval_dataset=val_data,
+    args=TrainingArguments(
+        num_train_epochs=1,
+        per_device_train_batch_size=1,
+        per_device_eval_batch_size=1,
+        learning_rate=3e-4,
+        gradient_accumulation_steps=4,
+        evaluation_strategy="steps" if VAL_SET_SIZE > 0 else "no",
+        save_strategy="steps",
+        eval_steps=2000 if VAL_SET_SIZE > 0 else None,
+        save_steps=2000,
+        output_dir=OUTPUT_DIR,
+        report_to = "tensorboard",
+        save_total_limit=3,
+        load_best_model_at_end=True if VAL_SET_SIZE > 0 else False,
+        optim="adamw_torch"
+    ),
+    data_collator=transformers.DataCollatorForSeq2Seq(tokenizer,
+                                                      pad_to_multiple_of=8,
+                                                      return_tensors="pt",
+                                                      padding=True),
+)
+
+
+trainer.train(resume_from_checkpoint=False)
+model.save_pretrained(OUTPUT_DIR)