增量预训练网络安全大模型的一次尝试

一、背景

探索使用网络安全知识，对开源基模型进行增强，评估是否能使基模型在网络安全领域表现出更好地专业度。

项目基于云起无垠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依然暂时是一个最佳实践。
• 安全领域数据集的定义是一个值得思考的问题，安全博客、安全论坛、安全论文这些语料就足够了吗？安全中特定领域任务的解决往往需要很多隐式知识以及隐式推理的能力，这些隐式知识所需要的数据在目前往往是十分稀缺和匮乏的，如果想要通过预训练得到一个能够解决大部分领域任务的基模型，数据工作可能十分任重道远。