Lion:专有大语言模型的对抗性蒸馏（EMNLP 2023）

[📄 论文] | [🤗 Lion权重]

新闻

[2023年10月8日] 我们的论文已被EMNLP 2023接收。
[2023年6月10日] 我们发布了解决微调过程中OOM问题的说明，请在训练过程中查看。
[2023年5月26日] 我们发布了模型权重。查看7B模型！
[2023年5月25日] 我们发布了在线演示，在这里试用我们的模型！
[2023年5月23日] 我们发布了用于训练和推理的代码。

概述

我们的对抗性蒸馏框架的高级概述，我们基于一个优秀的闭源LLM构建了一个紧凑的学生LLM，该闭源LLM扮演三个角色：教师、裁判和生成器。从左到右，一次迭代中有三个阶段：

_模仿_阶段，使学生的回应与教师的回应保持一致；
_区分_阶段，识别难样本；
_生成_阶段，生成新的难样本以增加学生模型面临的挑战。

恢复Lion权重

我们发布Lion权重作为增量权重，以遵守LLaMA模型许可。

Lion-7B（增量权重）您可以将我们的 delta 添加到原始 LLaMA 权重中以获得 Lion 权重。操作步骤：

按照此处的说明获取 huggingface 格式的原始 LLaMA 权重
请从 Hugging Face 下载我们的 delta 模型
使用以下脚本应用我们的 delta 来获取 Lion 权重：

python src/weight_diff.py recover --path_raw huggyllama/llama-7b --path_diff YuxinJiang/lion-7b --path_tuned <存储恢复权重的路径>

推理

对于 Lion 的推理和训练，请首先安装依赖：

pip install -r requirements.txt

我们提供了 Lion 的解码脚本，它读取输入文件并为每个样本生成相应的回复，最后将它们整合到一个输出文件中。这可以在单台配备 16GB GPU 的机器上运行。

python src/lion_inference.py \
    --model_dir <转换为hf格式的lion检查点和分词器路径> \
    --data_dir <输入json文件路径> \
    --output_dir <输出json文件路径> \
    --num_gpus 1

训练过程

以下展示了我们对抗性蒸馏框架的一次迭代。

1. 模仿阶段

1.1 获取教师模型对训练池的回应

python src/chatgpt_inference.py \
    -q <训练池json文件路径> \
    -o <训练池chatgpt推理结果路径> \
    --api_key <您的openai api密钥>

1.2 基于教师模型对训练池的回应对学生模型进行指令微调

微调在配备 8 张 A100 80G GPU 的机器上进行。

torchrun --nproc_per_node=8 --master_port=<您的随机端口> src/train.py \
    --model_name_or_path <转换为hf格式的检查点和分词器路径> \
    --data_path <训练池chatgpt推理结果路径> \
    --bf16 True \
    --output_dir result \
    --num_train_epochs 3 \
    --model_max_length 1024 \
    --per_device_train_batch_size 2 \
    --per_device_eval_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 600 \
    --save_total_limit 1 \
    --learning_rate 2e-5 \
    --weight_decay 0. \
    --warmup_ratio 0.03 \
    --lr_scheduler_type "cosine" \
    --logging_steps 1 \
    --fsdp "full_shard auto_wrap" \
    --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
    --tf32 True

解决内存不足问题

通常，微调 7B 模型需要约 7 x 8 x 2 = 112 GB 的显存。上述命令启用了参数分片，因此任何 GPU 上都不存储冗余的模型副本。如果您想进一步减少内存占用，以下是一些选项：

使用 --fsdp "full_shard auto_wrap offload" 为 FSDP 开启 CPU 卸载。这会以延长运行时间为代价节省显存。

根据我们的经验，DeepSpeed 阶段-3（带卸载）有时比带卸载的 FSDP 更节省内存。以下是使用 8 个 GPU 的 DeepSpeed 阶段-3 的示例，同时进行参数和优化器卸载：

deepspeed src/train_deepspeed.py \
    --model_name_or_path <转换为hf格式的检查点和分词器路径> \
    --data_path <训练池chatgpt推理结果路径> \
    --output_dir result \
    --num_train_epochs 3 \
    --model_max_length 1024 \
    --per_device_train_batch_size 16 \
    --per_device_eval_batch_size 1 \
    --gradient_accumulation_steps 1 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 600 \
    --save_total_limit 1 \
    --learning_rate 2e-5 \
    --warmup_ratio 0.03 \
    --logging_steps 1 \
    --lr_scheduler_type "cosine" \
    --report_to "tensorboard" \
    --gradient_checkpointing True \
    --deepspeed srcs/configs/deepspeed_config.json \
    --fp16 True

DeepSpeed 库还提供了一些有用的函数来估算内存使用情况。

LoRA 微调查询、键和值嵌入头的低秩切片。这可以将总内存占用从 112GB 减少到约 7x4=28GB。我们可能会在未来发布我们的重新实现，但目前 peft 代码库可作为有用的参考资源。

2. 区分阶段

2.1 获取教师模型对缓存池的回应

python src/chatgpt_inference.py \
    -q <缓存池json文件路径> \
    -o <缓存池chatgpt推理结果路径> \
    --api_key <您的openai api密钥>

2.2 获取学生模型对缓存池的回应

python src/lion_inference.py \
    --model_dir <HF转换的lion检查点和分词器的路径> \
    --data_dir <缓存池的JSON文件路径> \
    --output_dir <缓存池的lion推理输出路径> \
    --num_gpus 8

2.3 让裁判根据老师和学生的回答质量输出两个分数

为了减轻LLM裁判的位置偏差，我们通过交换老师回答和学生回答的位置进行两次运行。

python src/chatgpt_referee.py \
    -a <缓存池的chatgpt推理路径> <缓存池的lion推理路径> \
    -o <输出的chatgpt-lion评审文件路径> \
    --api_key <你的OpenAI API密钥>

python src/chatgpt_referee.py \
    -a <缓存池的lion推理路径> <缓存池的chatgpt推理路径> \
    -o <输出的lion-chatgpt评审文件路径> \
    --api_key <你的OpenAI API密钥>

2.4 区分困难指令和简单指令

python src/discrimination.py \
    --review12_path <输出的chatgpt-lion评审文件路径> \
    --review21_path <输出的lion-chatgpt评审文件路径> \
    --chatgpt_inference_path <缓存池的chatgpt推理路径> \
    --lion_inference_path <缓存池的lion推理路径> \
    --hard_save_path <识别出的困难指令保存路径> \
    --easy_save_path <识别出的简单指令保存路径>

3. 生成阶段

3.1 生成新的困难指令

python -m src/generate_hard_instruction generate_instruction_following_data \
    --seed_tasks_path <识别出的困难指令路径> \
    --all_tasks_path <缓存池的JSON文件路径> \
    --output_dir <生成的困难指令输出路径> \
    --num_instructions_to_generate 3000 \
    --api_key <你的OpenAI API密钥>

3.2 生成新的简单指令

python -m src/generate_easy_instruction generate_instruction_following_data \
    --seed_tasks_path <识别出的简单指令路径> \
    --all_tasks_path <缓存池的JSON文件路径> \
    --output_dir <生成的简单指令输出路径> \
    --num_instructions_to_generate 3000 \
    --api_key <你的OpenAI API密钥>

评估

开放式生成数据集的结果

我们利用GPT-4自动评估参考模型（ChatGPT）和候选模型之间回答的质量（评分范围1到10）。随后，我们将候选模型的表现计算为其相对于参考模型总分的百分比。

推理数据集的结果

引用

如果您使用了本仓库中的代码，请引用我们的论文。

@inproceedings{jiang-etal-2023-lion,
    title = "Lion: Adversarial Distillation of Proprietary Large Language Models",
    author = "Jiang, Yuxin  and
      Chan, Chunkit  and
      Chen, Mingyang  and
      Wang, Wei",
    editor = "Bouamor, Houda  and
      Pino, Juan  and
      Bali, Kalika",
    booktitle = "Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing",
    month = dec,
    year = "2023",
    address = "Singapore",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2023.emnlp-main.189",
    doi = "10.18653/v1/2023.emnlp-main.189",
    pages = "3134--3154",
}