基于KubeFATE的FATE-LLM任务实战

随着大型语言模型的不断蓬勃发展，相关新模型，新应用和新范式也在不断涌现，自 4 月发布以来，FATE-LLM 已经迭代发布了多个版本，不断完善大语言模型在联邦学习场景下的支持，以解决构建、使用大模型时的数据隐私保护问题以及公域数据短缺，可用数据不足的问题，在社区中得到了广泛的关注。

由 VMware AI Labs 团队发起并贡献的 KubeFATE 项目也在最近的多个版本中增强了对 FATE-LLM 在云原生环境下的支持，特别是针对 FATE-LLM 任务的专有需求，KubeFATE 在包括容器构建、GPU 调度与使用、模型存储等方面加入了专门的设计。本篇文章将给出一个基于 KubeFATE v1.11.2 和 FATE-LLM v1.2 的联邦大模型微调任务实例，并从任务设定、环境部署、所用技术、实验结果与分析等角度进行一个全方位的完整介绍。

任务设定与环境部署

在本文中，我们以 FATE-LLM v1.2 的 GPT2 教程为基础，设定一个两方的横向联邦学习场景，两方 Party ID 分别为 9999 和 10000，我们使用一个文本情感分类的任务来微调一个预训练的 GPT2 模型。该任务使用的数据集为 IMDB 影评数据集，与原教程的少量数据的示例不同，我们将训练数据集中全部 25000 条记录平均分给两方作为各方的训练数据，同时使用测试数据集的 25000 条记录作为评估数据。

对于实验环境，我们创建了两个 Kubernetes 集群，分别代表联邦任务的两个参与方。每个 Kubernetes 集群都包含一个 vSphere 虚拟机作为 GPU 节点，各虚拟机使用 PCI Passthrough 技术分别与一块 Nvidia V100 GPU 集成。我们使用 Docker 和 cri-docker，以及 nvidia-container-runtime 和 k8s-device-plugin 等组件以在 Kubernetes 集群中使用该 GPU。本实验所有关键依赖项的具体版本如下：

我们可以按照 KubeFATE 项目 GitHub 仓库中的 K8s 环境使用指南来部署 KubeFATE 和 FATE 集群。要使用 FATE-LLM，我们需要在用于部署 FATE 集群的 cluster.yaml 文件中显式指定某些设置。首先我们应该将 algorithm 参数设置为 “ALL”，并将 device 参数设置为 “GPU”，这表示我们要使用包含 FATE-LLM 和相关模块，以及带有 GPU 驱动和库的 FATE 容器镜像。此外，我们需要在 python 组件即 FATE-Flow 容器的配置中，在 resources 资源部分加入 GPU 资源的请求，例如本文使用了 “nvidia.com/gpu: 1”。以下是这些设置的示例：

需要注意的是，本文示例中 FATE-LLM 的训练任务是由 FATE-Flow 容器执行的，因此 GPU 资源分配给了名为 python 的容器。对于使用了 DeepSpeed 的 FATE-LLM 任务，我们需要为 nodemanager 组件配置该 GPU 资源设定。此外，我们建议为 FATE 集群开启持久化，从而使 KubeFATE 能够保存 FATE-LLM 任务中的预训练模型、微调模型、任务记录、日志等，防止这些关键文件在容器发生重启后被清理。

为了验证 FATE 集群部署后的环境和设置，我们可以使用 kubectl exec 进入 fateflow pod 并使用 nvidia-smi 命令检查其中可用的 GPU 资源。在 FATE 集群部署好后，我们就可以开始发起 FATE-LLM 任务了，FATE-LLM v1.2 中使用 HuggingFace 生态的 peft 库来支持多种高效的参数微调方法，在本文的实践中，我们将采用 LoRA、Prefix Tuning、Prompt Tuning 和 P-Tuning 的方法进行比较实验。

采用的 PEFT 方法

为了使用大语言模型来完成特定的下游任务，如果需要对模型的全部参数进行微调，会造成了巨大的时间开销与数据存储、传输成本。PEFT 方法在保持模型性能相当的前提下，通过减少微调的参数量以降低计算、存储、传输数据的成本。在本节中，我们将简要介绍采用的四种 PEFT 方法。

• LoRA（Low-Rank Adaptation）：LoRA 的核心思想是在适应新任务时，权重矩阵的更新不必与原始权重矩阵具有相同的 “秩”。因此对每个权重矩阵，我们可以引入两个低秩矩阵模块去代表该更新，微调时则仅训练更新这些低秩的矩阵。在实践中，LoRA 一般被应用到 attention 模块中。
• Prefix Tuning：该方法在每个 transformer 层的输入之前构造一段任务相关的 virtual tokens 作为 prefix，在训练时只更新 prefix 部分的参数，而固定 transformer 中的其他部分参数。此外，为了防止直接更新 prefix 的参数导致训练不稳定的情况，prefix 层前可以加入了 MLP（Multi-layer Perception）结构，即将 prefix 分解为更小维度的 input 与 MLP 的组合。
• Prompt Tuning：该方法可以看作是 Prefix Tuning 的简化版本，它只在输入层加入 prompt tokens，而并不需要加入 MLP 进行调整。
• P-Tuning：该方法的思想也是在输入层添加可训练的 prompt 参数，由于 prompt 的词之间是彼此关联的，离散的 prompt 对于连续的神经网络并不是最优解，我们需要采用某种方法将它们关联起来。于是 P-Tuning 将一些伪 prompt 输入至 LSTM 或类似模型中，利用 LSTM 的输出向量来替代原始的 prompt token，最后一起输入至预训练语言模型中。

运行 FATE-LLM 任务

接下来我们便可以着手运行 FATE-LLM 任务。首先，guest 与 host 两方需将本地数据上传至 FATE 系统，即将预处理后的 CSV 文件分别放入各方的 FATE-Flow 容器中。然后我们可以在 KubeFATE 提供的 Jupyter Notebook 中使用以下代码将上传的数据绑定到 FATE 中。（对于 host 方，需要绑定 data_host 数据）

随后我们便可以在 guest 方提交联邦学习任务，具体流程与 FATE-LLM 仓库中的 GPT2 微调教程基本一致，以下是需额外注意的几点：

• 因为本文加入了用于评估的测试数据，我们可以新建一个 Reader 组件读取该数据并作为 NN 组件的 validate_date；
• 若任务出现超时或 Pin memory thread 退出等异常，可以尝试将 “save_to_local_dir=True, pin_memory=False” 添加到 TrainerParam；
• 对于本文使用的二分类任务，我们为其配置了 Evaluation 组件用于评估训练后的模型的性能。

我们使用了多种不同的参数设置方式运行了若干个 FATE-LLM 任务来评估它们对训练过程与模型性能的影响：

• 设置 TrainerParam 中的 CUDA 参数来比较训练时使用 CPU 和 GPU 的差异
• 设定 peft_type 与 peft_config 参数来采用不同的 PEFT 方法以及配置
• 设置 aggregate_every_n_epoch 参数比较不同的本地训练轮次对训练效果的影响
• 调整 TrainerParam 中的 epoch 参数找到模型效果最好的轮次（最大为 10）
• 我们实现了非 PEFT 的 GPT2 Full Fine-tuning 方法，并以此作为对照进行了对比实验

实验结果分析

我们对使用的设备类型、PEFT 方法以及运行的 epoch 次数均进行了多元化的设置。模型性能通过 AUC、F1-score 等机器学习指标进行评估。除此以外，我们还列举了若干实验本身的数据，包括聚合过程的数据传输量，任务运行时间等，便于进行更深入的分析。在使用 PEFT 方法对模型训练过程进行优化时，我们尽量使用 peft 库默认的参数配置。以下是我们实验的部分结果数据：

通过对比表格中与模型性能相关的指标，我们可以看到 LoRA、Prefix Tuning 方法与 Full Fine-tuning 对照实验的模型性能相当。与此同时，它们的数据传输量、整体训练时间也显著低于 Full Fine-tuning，更进一步，我们观察到设定更大一点的模型聚合间的本地训练轮次（同时总轮次不变），能够进一步降低数据传输成本，同时不影响模型性能，这证实了联邦学习与 PEFT 方法之于微调大语言模型是可行且高效的。而对于另两种 PEFT 方法，模型性能则有较大的差距，我们猜测这与可训练参数过少或相关 PEFT 组件参数设置有关，导致模型泛化能力受限。我们在接下来的实验中也将尝试其他方法和参数组合以进一步验证。

总结与展望

本文介绍了使用 KubeFATE 从部署、配置到发起运行 FATE-LLM 任务的完整流程。我们展示了不同配置下的实验结果，并进行了分析。实验数据证实了将大语言模型与联邦学习架构结合的合理性，并体现出 PEFT 方法的强大性能。其在显著降低通信成本、数据存储成本的同时，将模型性能维持在一个相当可观的水平。

本文提供的实例基于 FATE-LLM v1.2 版本，可作为横向联邦场景下，使用云原生基础设施对同构的大语言模型进行 PEFT 的微调训练的一个参考。除 GPT2 示例以外，KubeFATE 也支持 FATE-LLM v1.2 中的其他模型和训练方法，例如使用 DeepSpeed 进行 ChatGLM-6B、Llama 等模型的训练。而在最新发布的 v1.3 版本中，FATE-LLM 项目引入了 FTL-LLM 面向大语言模型的联邦迁移学习范式，并实现了联邦 Offsite-Tuning 框架。我们会在后续的 KubeFATE 迭代中进一步加强相关的支持，也欢迎社区开发者和用户的参与和关注。

内容来源｜公众号：VMware 中国研发中心

有任何疑问，欢迎扫描下方公众号联系我们哦～