OpenFedLLM：通过联邦学习在分布式私有数据上训练大型语言模型

OpenFedLLM是一个用于通过联邦学习训练*大型语言模型（LLM）*的开源研究代码库。详情请查看我们的论文及相应的实证研究。

OpenFedLLM包含以下主要特点：

7种联邦学习算法（如FedAvg、FedProx、SCAFFOLD、FedAvgM等）。
2种LLM训练算法，包括指令微调（即SFT）和价值对齐（即DPO）。
30多种评估指标，涵盖通用能力、医疗问答、金融问答、代码生成、数学解题等方面。

新闻🔥

2024年6月： 我们发布了首个现实的FedLLM基准测试：FedLLM-Bench。查看论文 | 代码。

安装

克隆仓库、子模块并安装所需包。

git clone --recursive --shallow-submodules https://github.com/rui-ye/OpenFedLLM.git
cd OpenFedLLM
conda create -n fedllm python=3.10
conda activate fedllm
pip install -r requirements.txt
source setup.sh

训练

我们在training_scripts/下提供了训练脚本。请在此仓库的顶层目录中尝试运行它们。

联邦指令微调

训练脚本位于training_scripts/run_sft.sh。

CUDA_VISIBLE_DEVICES=1 python main_sft.py \
 --model_name_or_path "meta-llama/Llama-2-7b-hf" \
 --dataset_name "vicgalle/alpaca-gpt4" \
 --dataset_sample 20000 \
 --fed_alg "fedavg" \
 --num_clients 20 \
 --sample_clients 2 \
 --max_steps 10 \
 --num_rounds 200 \
 --batch_size 16 \
 --gradient_accumulation_steps 1 \
 --seq_length 512 \
 --peft_lora_r 32 \
 --peft_lora_alpha 64 \
 --use_peft \
 --load_in_8bit \
 --output_dir "./output" \
 --template "alpaca" \

主要参数：

model_name_or_path：基础模型的名称或本地位置
template：对话模板。可在utils/template.py中定义自己的模板。
dataset_name：数据集名称。如果您感兴趣的数据集尚未支持，可以修改utils/process_dataset.py。
dataset_sample：如需从原始数据集中采样特定数量的样本，则需要此参数。
fed_alg：联邦学习算法的名称
num_clients/sample_clients：总共num_clients个客户端，每轮选择sample_clients个客户端
max_steps：每轮每个客户端的模型更新步数。

联邦价值对齐

训练脚本位于training_scripts/run_dpo.sh。

python main_dpo.py --template "vicuna_v1.1"

注意，main_sft.py和main_dpo.py在使用上的主要区别在于template参数。我们计划在未来使它们保持一致。

对于SFT，模板定义在utils/template.py中
对于DPO，模板定义在utils/conversation.py中

评估

评估代码放在evaluation/目录下。我们的大多数评估遵循现有的高影响力开源仓库。请参考各子目录中的详细README和运行脚本。

例如，evaluation/open_ended/包含了三个基准测试的开放式评估，涵盖MT-Bench、Vicuna Bench和AdvBench；详见README.md。

引用

如果您发现本仓库有帮助，请引用我们的论文。

@article{ye2024openfedllm,
  title={OpenFedLLM: Training Large Language Models on Decentralized Private Data via Federated Learning},
  author={Ye, Rui and Wang, Wenhao and Chai, Jingyi and Li, Dihan and Li, Zexi and Xu, Yinda and Du, Yaxin and Wang, Yanfeng and Chen, Siheng},
  journal={arXiv preprint arXiv:2402.06954},
  year={2024}
}