LLM大模型GPT2微調嘗試

第七子007發表於2024-05-20
大模型GPT

　 1、作為安全從業者，以前搞逆向、挖漏洞、幹滲透全靠人工推進，缺點很明顯：

  • 無法自動化，甚至也無法半自動化，效率低（後續可以開發agent解決）
  • 知識面有限，存在很多知識盲點，導致遇到部分問題無法解決（可以透過增加知識庫，然後rag檢索或微調大模型解決）

嘗試了一些線上的大模型(chatGPT4、copilot)，挖掘 https://www.cnblogs.com/theseventhson/p/13933230.html 這個例子的棧溢位漏洞，效果還行, 能找到漏洞；離線部署的LLM中嘗試了secGPT、openbuddy等，同樣的程式碼也能準確找到漏洞所在，比如：

說明基於現有大資料訓練得到的模型是能夠發現並解決已知問題的，這條路別人已經走通了！

2、 訓練語料

（1）眾所周知，機器學習分監督學習於非監督學習。不論哪種，訓練資料都是必須的，所以先整理一下我自己的部落格園技術總結：

huggingface上大模型sft微調一般的資料格式為：

{
    "instruction": "Please perform static application security testing on the above code",
    "input": "static void goodG2B1()\n{\n    int data;\n    /* Initialize data */\n    data = -1;\n    if(0)\n    {\n        /* INCIDENTAL: CWE 561 Dead Code, the code below will never run */\n        printLine(\"Benign, fixed string\");\n    }\n    else\n    {\n        /* FIX: Set data to a relatively small number greater than zero */\n        data = 20;\n    }\n    {\n        size_t i;\n        int *intPointer;\n        /* POTENTIAL FLAW: if data * sizeof(int) > SIZE_MAX, overflows to a small value\n         * so that the for loop doing the initialization causes a buffer overflow */\n        intPointer = (int*)malloc(data * sizeof(int));\n        if (intPointer == NULL) {exit(-1);}\n        for (i = 0; i < (size_t)data; i++)\n        {\n            intPointer[i] = 0; /* Potentially writes beyond the boundary of intPointer */\n        }\n        printIntLine(intPointer[0]);\n        free(intPointer);\n    }\n}\n",
    "output": "this code does not violate any security encoding standards"
}

　 這類資料需要標註，耗時耗力，而部落格園的技術文章都是沒有標註的，所以sft暫時不考慮！如果非要把部落格園的文章搞成instrct格式，可以嘗試把文章內容作為input，標題或主要內容作為output，instruction改成自己關注的問題就好！

（2）我因為沒時間標註，直接把文章內容挨個放入json檔案；為了減少訓練時記憶體佔用，也為了GPT2的max_length=1024的限制，把文章按照1024的長度截斷，分別存入json檔案；一共準備了19個json檔案；

用於微調的樣本資料：

3、模型選擇：huggingface上模型已多大66w，這麼多模型怎麼選合適的？

（1）模型引數量：N = l * 12* d^2; l是transformer block的個數；d 是 embedding的維度；

（2）模型總的計算量：C=6*ND； N是模型引數量，D是訓練集的token總數

（3）模型BP需要儲存：：16byte*N + FFN(d、樣本length、batch_size等)

由於token數量只有200w，如果選擇引數大的模型，肯定欠擬合（參考scaling laws，我這點token量都達不到人家嘗試資料量的下限），所以只能選擇引數小的模型；這裡最終選擇GPT2嘗試！

4、微調的方式有很多種，這裡選擇截至目前最優的lora嘗試：

  1 import logging
  2 import torch
  3 from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments
  4 from datasets import load_dataset
  5 from peft import get_peft_model, LoraConfig, TaskType
  6 import os
  7 
  8 logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)
  9 
 10
 11 def preprocess_function(examples):
 12     inputs = tokenizer(examples["text"], truncation=True, padding="max_length", max_length=1024)
 13     inputs["labels"] = inputs["input_ids"].copy()
 14     return inputs
 15 
 16 if __name__ == '__main__':
 17
 18     model_name = "gpt2"
 19     tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 20     model = GPT2LMHeadModel.from_pretrained(model_name)
 21     logging.info("Model loaded successfully")
 22 
 23     
 24     tokenizer.pad_token = tokenizer.eos_token
 25 
 26     
 27     training_args = TrainingArguments(
 28         output_dir="/root/huggingface/GPT2/Lora",
 29         overwrite_output_dir=True,
 30         num_train_epochs=1,
 31         per_device_train_batch_size=20,
 32         save_steps=10_000,
 33         save_total_limit=2,
 34         logging_dir="/root/huggingface/GPT2/logs",
 35         save_strategy="epoch",  
 36     )
 37 
 38     lora_config = LoraConfig(
 39         task_type=TaskType.CAUSAL_LM,
 40         r=16,
 41         lora_alpha=32,
 42         lora_dropout=0.1,
 43         target_modules=["attn.c_proj", "attn.c_attn"]
 44     )
 45 
 46     
 47     model = get_peft_model(model, lora_config)
 48     model.print_trainable_parameters()  # 列印可訓練引數
 49 
 50     last_checkpoint = None
 51     checkpoint_prefix = "checkpoint"
 52 
 53     # 檢查是否存在之前的檢查點
 54     for i in range(19, 0, -1):
 55         checkpoint_dir = f"/root/huggingface/GPT2/{checkpoint_prefix}-{i}"
 56         if os.path.exists(checkpoint_dir):
 57             last_checkpoint = checkpoint_dir
 58             logging.info(f"last_checkpoint:{last_checkpoint}")
 59             break
 60 
 61     logging.info(f"last_checkpoint:{last_checkpoint}")
 62     # 從上一個檢查點繼續訓練
 63     start_file_index = 1 if last_checkpoint is None else int(last_checkpoint.split('-')[-1]) + 1
 64 
 65     # 遍歷每個資料檔案並進行訓練
 66     for i in range(start_file_index, 20):
 67         data_file = f'/root/huggingface/data/cyber_security{i}.json'
 68         dataset = load_dataset('json', data_files=data_file, split='train')
 69         logging.info(f"Dataset {data_file} loaded successfully")
 70         print(f"Dataset size: {len(dataset)}")
 71         tokenized_dataset = dataset.map(preprocess_function, batched=True)
 72         print(f"Tokenized dataset size: {len(tokenized_dataset)}")
 73 
 74         # 建立Trainer物件
 75         trainer = Trainer(
 76             model=model,
 77             args=training_args,
 78             train_dataset=tokenized_dataset,
 79         )
 80 
 81         logging.info(f"Before training on {data_file}")
 82 
 83         if last_checkpoint:
 84             logging.info(f"Loading from checkpoint {last_checkpoint}")
 85             trainer.train(resume_from_checkpoint=last_checkpoint)
 86             last_checkpoint = None  # 只在第一次載入檢查點
 87         else:
 88             # 訓練模型
 89             trainer.train()
 90 
 91         logging.info(f"After training on {data_file}")
 92 
 93         # 儲存中間模型和檢查點
 94         model.save_pretrained(f"/root/huggingface/GPT2/fine-tuned-model-{i}")
 95         tokenizer.save_pretrained(f"/root/huggingface/GPT2/fine-tuned-model-{i}")
 96         trainer.save_model(output_dir=f"/root/huggingface/GPT2/{checkpoint_prefix}-{i}")
 97         logging.info(f"Model saved successfully after training on {data_file}")
 98 
 99     # 最終儲存模型
100     model.save_pretrained("/root/huggingface/GPT2/fine-tuned-model-final")
101     tokenizer.save_pretrained("/root/huggingface/GPT2/fine-tuned-model-final")
102     logging.info("Final model saved successfully")

　　為了測試不同的引數的效果，我這裡分別使用了r=8和r=16兩個引數分別測試：

r=8的trainable：

cross entropy的loss：

r=16的trainable：

cross entropy的loss：

微調完成後推理：中文效果非常差，感覺完全沒理解語義

英語效果也好不到哪去：

原因：

  • GPT2訓練的中文語料少，對中文支援不足
  • 微調的語料太少，每批訓練語料的epoch也只有1，loss最小也在3；

總結：

1、transformer架構，基礎的block塊核心由兩部分構成：attention和FFN；FFN就是深度神經網路，十幾年前就有了，沒啥新的，所以attention是較大創新；

  • 在attention之前的embedding加上了位置編碼，完美記錄了token的位置，更好地表達語義；
  • attention本身透過q和k的向量內積提取token之間的相似度，找到每個token重要的context
  • q*k的乘積作為權重調整v的值，更好地表達token在當前context中的值；

比如“apple”這個token，初始的embedding都是一樣的，無法區分是水果的apple，還是電子產品的apple，只能根據不同的context確定apple這個詞的語義；如果apple周圍出現大量的pear、bananer、peach等，大機率指的是水果apple，理論上講multihead中水果類head轉換的v值會被q*k的內積增加，而電子產品類head轉換的v值會被q*k的內積減小；所以，attention最核心的功能：根據context調整token的v值，讓同一token在不同的context有不同的v值，更利於後續進一步的語義理解！目的感覺和網際網路搜廣推的“千人千面”很像啊！所以使用模型時處理token長度的max_length值不能太小（普通問答1024夠了，專業的技術文章建議至少10240），否則提取的context資訊有限，影響token v值的正確調整！

2、transformer每個block都幹同樣的事：

  • attention：根據context調整token的v值
  • FFN：空間轉換提取特徵

理論上講：只要訓練語料足夠多，block數量足夠多，經過層層變換，總能精準得到每個token的v值和語義資訊，所以transformer的核心是“大力出奇跡”！

附：

1、知識庫rag和微調的優劣對比：

相關文章