基于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 中国研发中心

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

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章