增量預訓練網絡安全大模型的一次嘗試

一、背景

探索使用網絡安全知識，對開源基模型進行增強，評估是否能使基模型在網絡安全領域表現出更好地專業度。

項目基於雲起無垠SecGPT開源項目，在hugeface開源數據集的基礎上，增加了自有預訓練數據，進行增量預訓練。

參考鏈接：

https://github.com/Clouditera/secgpt

 

二、數據集準備

0x1：預訓練數據

預訓練使用“安全書籍，安全知識庫，安全論文，安全社區文章，漏洞庫”等等安全內容，

https://huggingface.co/datasets/w8ay/security-paper-datasets?row=0 

筆者增加了一份自己近10年以內的博客文章，

將上述格式轉換爲hugeface兼容格式，huggingface支持以下4種數據格式的數據集，只需要在load的時候設定格式就好了。

Data format	Loading script	Example
CSV & TSV	csv	load_dataset("csv", data_files="my_file.csv")
Text files	text	load_dataset("text", data_files="my_file.txt")
JSON & JSON Lines	json	load_dataset("json", data_files="my_file.jsonl")
Pickled DataFrames	pandas	load_dataset("pandas", data_files="my_dataframe.pkl")

# -- coding: utf-8 --**

import json
import pandas as pd
from datasets import load_dataset
from datasets import Dataset, concatenate_datasets


def prepare_cnblogs_data():
    with open("./posts.json", 'r', encoding='utf-8') as file:
        data = json.load(file)

    original_json = list()
    for item in data:
        original_json.append(
            {
                "text": item['Body'],
                "category": "cnblogs"
            }
        )

    print(original_json[0])
    print(original_json[1])

    with open("./posts_hugeface_dataset.json", 'w', encoding='utf-8') as file:
        # 將列表保存到文件，注意使用ensure_ascii=False參數
        json.dump(original_json, file, ensure_ascii=False, indent=2)

    df = pd.DataFrame(original_json)
    df.to_pickle('./posts_hugeface_dataset.pkl')


def prepare_security_paper_dataset():
    dataset = load_dataset("w8ay/security-paper-datasets")
    df_train = dataset['train'].to_pandas()
    df_train.to_pickle('security_paper_dataset.pkl')


def merge_dataset():
    security_paper_dataset = load_dataset("pandas", data_files="security_paper_dataset.pkl")['train'].to_pandas()
    posts_hugeface_dataset = load_dataset("pandas", data_files="posts_hugeface_dataset.pkl")['train'].to_pandas()
    dataset_part1 = Dataset.from_pandas(security_paper_dataset)
    dataset_part2 = Dataset.from_pandas(posts_hugeface_dataset)
    combined_dataset = concatenate_datasets([dataset_part1, dataset_part2])

    print(combined_dataset[0])
    print(len(combined_dataset))

    combined_dataset_df = pd.DataFrame(combined_dataset)
    combined_dataset_df.to_pickle('security_paper_and_posts_dataset.pickle')


if __name__ == "__main__":
    prepare_security_paper_dataset()
    prepare_cnblogs_data()
    merge_dataset()

這裏將多份本地pickle語料數據合併，

0x2：有監督數據

通過構造各類有監督安全能力數據集，讓模型能瞭解各類安全指令。 

有監督SFT數據集構建有幾個注意點：

• 思維鏈：基於思維鏈方式構造有監督數據集讓模型能夠根據問題逐步推理到最終答案，展現推理過程。
• 爲防止災難性遺忘，有監督數據喂通用能力數據+安全能力數據，數據佔比5:1

參考鏈接：

https://huggingface.co/w8ay/secgpt
https://huggingface.co/datasets/TigerResearch/tigerbot-zhihu-zh-10k
https://zhuanlan.zhihu.com/p/548355568
https://zhuanlan.zhihu.com/p/564816807
https://zhuanlan.zhihu.com/p/554678463
https://huggingface.co/datasets/w8ay/security-paper-datasets
https://huggingface.co/w8ay/secgpt

 

三、模型訓練

0x1：基模型選擇

基於Baichuan-13B (無道德限制，較好中文支持，顯存資源佔用小)

Baichuan-13B 是由百川智能繼 Baichuan-7B 之後開發的包含 130 億參數的開源可商用的大規模語言模型，在權威的中文和英文 benchmark 上均取得同尺寸最好的效果。本次發佈包含有預訓練 (Baichuan-13B-Base) 和對齊 (Baichuan-13B-Chat) 兩個版本。Baichuan-13B 有如下幾個特點：

• 更大尺寸、更多數據：Baichuan-13B 在 Baichuan-7B 的基礎上進一步擴大參數量到 130 億，並且在高質量的語料上訓練了 1.4 萬億 tokens，超過 LLaMA-13B 40%，是當前開源 13B 尺寸下訓練數據量最多的模型。支持中英雙語，使用 ALiBi 位置編碼，上下文窗口長度爲 4096。
• 同時開源預訓練和對齊模型：預訓練模型是適用開發者的“基座”，而廣大普通用戶對有對話功能的對齊模型具有更強的需求。因此本次開源我們同時發佈了對齊模型（Baichuan-13B-Chat），具有很強的對話能力，開箱即用，幾行代碼即可簡單的部署。
• 更高效的推理：爲了支持更廣大用戶的使用，我們本次同時開源了 int8 和 int4 的量化版本，相對非量化版本在幾乎沒有效果損失的情況下大大降低了部署的機器資源門檻，可以部署在如 Nvidia 3090 這樣的消費級顯卡上。
• 開源免費可商用：Baichuan-13B 不僅對學術研究完全開放，開發者也僅需郵件申請並獲得官方商用許可後，即可以免費商用。

0x2：運行環境

V100 * 8 

0x3：增量預訓練 

啓動訓練，

python train.py

參考鏈接：

https://huggingface.co/baichuan-inc/Baichuan-13B-Base
https://huggingface.co/w8ay/secgpt
https://zhuanlan.zhihu.com/p/554678463

 

四、效果展示

模型結果的輸出存在隨機性，模型可能知道答案，但是隨機後改變輸出，這時候可以增加提示詞讓模型注意力集中。

也可以做RLHF強化學習，讓模型輸出多個結果，選擇正確的那個，提升模型效果。

0x1：用百川原生基模型

python3 webdemo.py --base_model baichuan-inc/Baichuan-13B-Base

0x2：用增量預訓練後的模型

 

python3 webdemo.py --base_model baichuan-inc/Baichuan-13B-Base --lora /data_vdb1/secgpt/output/epoch-0-step-2400

 

http://47.236.142.150:7860/

參考鏈接：

https://github.com/Clouditera/secgpt

 

五、後續思考

• 增量預訓練成本（時間/金錢）相對較大，8張V100跑了1天多才能跑完一個epoch，但ROI並不明顯，沒有經過RLHF的增量預訓練基模型，在筆者自己的安全領域問題基線集中，變現並不比SFT的效果要好。目前來看，在特定安全領域問題內，SFT依然暫時是一個最佳實踐。
• 安全領域數據集的定義是一個值得思考的問題，安全博客、安全論壇、安全論文這些語料就足夠了嗎？安全中特定領域任務的解決往往需要很多隱式知識以及隱式推理的能力，這些隱式知識所需要的數據在目前往往是十分稀缺和匱乏的，如果想要通過預訓練得到一個能夠解決大部分領域任務的基模型，數據工作可能十分任重道遠。