bloom_tuning: BLOOM 模型的指令微调

介绍

BLOOM 系列模型是由数百名研究人员在包含 46 种自然语言和 13 种编程语言的数据集上基于大规模分布式训练框架 Megatron-DeepSpeed 训练得到。实验发现，BLOOM 在一系列基准测试上取得了具有竞争力的性能，经过多任务提示微调后，可以获得更为惊艳的效果。

BLOOM 模型支持中文、英文、代码、法语、西班牙语。训练数据的具体语言占比如下：

链接：https://huggingface.co/bigscience/bloom-560m

LLMPruner 工具对 BLOOM 进行词表裁剪，保留常用的中英文 token，词表大小由 250880 降至 46145，缩减为原来的 18.39%，在后续微调过程中可以减少显存占用。

词表裁剪后的模型链接：https://huggingface.co/YeungNLP/bloom-396m-zh

数据

训练数据来自于 BelleGroup/train_3.5M_CN，该数据集包含 3.6M 条指令，从中筛选出单轮对话数据，进行10:1采样后得到约 0.25M 指令数据：

python sample_data.py \
--input data/train_3.5M_CN.json \
--output data/train.jsonl \
--sample_ratio 0.1

单条指令数据形如：

{
    "instruction": "你好，请问你能做什么？", 
    "output": "你好，我可以回答各种问题，提供辅助，或者与你聊天。有什么我可以帮你的吗？"
}

输出部分的长度分布如下图所示（若输出长度超过2048，则设置为2048）：

可以构建自己的指令数据集进行训练。

指令微调

单卡进行指令微调：

export CUDA_VISIBLE_DEVICES=0
python train.py \
--model_name_or_path /path/to/bloom \
--data_path data/train.jsonl \
--max_input_length 200 \
--max_output_length 768 \
--output_dir output \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 16 \
--learning_rate 3e-5 \
--num_train_epochs 2 \
--lr_scheduler_type "cosine" \
--warmup_steps 2000 \
--logging_steps 10 \
--save_strategy "steps" \
--save_steps 200 \
--save_total_limit 1 \
--fp16 False

基于 deepspeed ZeRO-Stage 2 (未启用 offload_optimizer - CPU) 进行指令微调训练：

deepspeed --include localhost:0 train.py \
--model_name_or_path /path/to/bloom \
--data_path data/train.jsonl \
--max_input_length 200 \
--max_output_length 768 \
--output_dir output \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 16 \
--learning_rate 3e-5 \
--num_train_epochs 2 \
--lr_scheduler_type "cosine" \
--warmup_steps 2000 \
--logging_steps 10 \
--save_strategy "steps" \
--save_steps 200 \
--save_total_limit 1 \
--deepspeed deepspeed.json \
--fp16 False

在 per_device_train_batch_size=1、gradient_accumulation_steps=16、max_input_length=200、max_output_length=768、fp16=false 的配置下，需要约 14G 显存对 bloom-396m-zh 进行微调。在显存不足的情况下，建议启用 offload_optimizer (Enable offloading of optimizer state and optimizer computation to CPU) 和设置 fp16=true 进一步降低显存占用。训练过程中的 loss 变化如下图所示：

推理

微调后的模型已上传至 huggingface: bloom-396m-chat

可以通过如下代码调用 bloom-396m-chat 模型来生成对话：

from transformers import BloomTokenizerFast, BloomForCausalLM

model_name_or_path = "WangZeJun/bloom-396m-chat"

tokenizer = BloomTokenizerFast.from_pretrained(model_name_or_path)
model = BloomForCausalLM.from_pretrained(model_name_or_path).cuda()
model = model.eval()

input_pattern = "{}</s>"
text = "你好"
input_ids = tokenizer(input_pattern.format(text), return_tensors="pt").input_ids
input_ids = input_ids.cuda()

outputs = model.generate(input_ids, do_sample=True, max_new_tokens=1024, top_p=0.85,
    temperature=0.3, repetition_penalty=1.2, eos_token_id=tokenizer.eos_token_id)

output = tokenizer.decode(outputs[0])
response = output.replace(text, "").replace('</s>', "")
print(response)

运行如下命令进行命令行推理：

python chat.py \
--model_name_or_path /path/to/model \
--max_new_tokens 1024 \
--top_p 0.85 \
--temperature 0.3 \
--repetition_penalty 1.2

部分示例如下：

引用

若使用本项目的代码或模型，请引用本项目。

@misc{bloom_tuning,
  author = {Zejun Wang},
  title = {bloom_tuning: BLOOM 模型的指令微调},
  year = {2023},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/zejunwang1/bloom_tuning}}
}