【隐私计算篇】使用GPU加速计算联邦学习XGBOOST算法以及对NVIDIA FLARE（NVIDIA 联邦学习应用运行环境）的介绍

1. 背景介绍

借着最近在搞GPU相关的项目契机，来介绍一下英伟达Nvidia FLARE项目【1】，并且利用GPU硬件来加速联邦学习XGBOOST算法。感觉开源的机器学习、深度学习已经开始出现拥抱隐私计算的趋势，比如近期我正在关注Andrew Ng的联邦学习用于大模型的微调课程【2】，利用开源横向联邦学习框架flower【3】进行实现，这个我近几天有时间会做一期关于联邦学习实现大模型微调的技术分享。另外，看到xgboost的原生代码仓库【4】，已经添加了federated-secure以及vertical-federated-learning联邦学习分支。英伟达也推出了Nvidia FLARE【5】联邦学习框架。

接下来会介绍Nvidia FLARE，预期收益是通过了解开源的联邦计算框架组成、思想、运行模式，进而了解联邦学习的基本算法思路，当然这里需要说明下联邦学习的框架有多种方案，FLARE只是其中一种模式。并利用Nvidia FLARE框架，来实现GPU加速联邦学习xgboost。

2. 联邦计算框架介绍

NVIDIA FLARE（NVIDIA 联邦学习应用运行环境）是一个可扩展SDK，能够将现有的机器学习/深度学习（ML/DL）和计算工作流适配为联邦范式。通过FLARE，能够支持创建一个面向隐私保护的去中心化的数据计算解决方案，促进分布式多方协作计算。

FLARE面向联邦计算，其核心是一个联邦计算框架，支持构建联邦学习和联邦分析等应用。与需要将数据复制到集中位置的集中数据湖解决方案不同，联邦计算框架是将计算能力带到分布式数据集，确保数据保留在计算节点本地，只有经过批准的结果（或者密文）在协作方之间共享。

2.1 开发工具

一个完备的联邦计算框架一般可以提供在开发过程不同阶段的多种工具，包含以下：

联邦计算框架客户端API：用户只需修改几行代码，即可从ML/DL切换到联邦学习。
模拟器CLI：允许用户在单台计算机中模拟多进程环境中的联邦学习或计算任务。
POC CLI：提供联邦网络的模拟，支持用户在单台主机上模拟项目部署。
Job CLI：允许用户直接在POC或生产环境中创建和提交任务。
联邦计算框架API：用户可以直接通过Python代码或notebooks运行任务。
联邦计算框架仪表盘：支持用户设置、批准和分发部署工件。
预检工具：帮助用户在运行任务之前验证联邦系统是否正确设置。
云部署CLI：通过一个CLI命令启动并部署联邦计算框架。
ML实验跟踪支持：支持TensorBoard、MLFlow和Weights & Biases的实验跟踪。

2.2 安全性和隐私保护能力

这个问题很多甲方（B端客户）其实非常关心，部署应用隐私计算平台，首先需要说明安全性是如何保护，进而需要证明安全性真的得到了保护。

安全配置：使用TLS证书确保安全环境。
事件驱动的安全插件机制：支持用户定义的本地身份验证和授权。
授权策略控制：本地实体可以控制联邦框架内的授权策略。
数据与结果过滤机制：增强数据保护。
审计日志：提供透明性和责任追踪。
区块链：支持全流程或者关键流程的行为记录上链。
隐私保护算法：包括同态加密、安全多方计算、安全求交、匿踪查询等，这些隐私保护技术可以参考我们的技术专栏《隐私计算》，也可以选择特定的几篇《中国剩余定理解释以及Paillier解密加速应用》、《MPC安全多方计算矩阵乘法算子的原理分析》、《隐私计算匿踪查询技术深入浅出》、《隐私集合求交（PSI）原理深入浅出》。

2.3 框架设计理念

每个组件和 API 都是基于规范的，可以通过遵循规范来构建替代实现。用户能够进行扩展和定制，以满足特定工作流的需求。能够处理真实世界中的用例，可以应对意外事件或异常代码。考虑系统的通用性，以支持不同的“联邦”计算用例。将组件谨慎地打包到不同的层次中，确保各层之间的依赖最小化。这种设计目前看其实非常适合实际需求，因为一开始隐私计算厂商，会将安全求交、匿踪查询、联邦学习、安全多方计算等不同的技术整合到一个整体来实现，造成系统的庞大和臃肿，但到了真正的业务端，其实可能只会用到其中的某一种技术或者部分技术，并不需要部署全部的服务，所以就涉及到了拆解和瘦身。

2.4 系统概览

NVFlare 实际上是一种支持多个客户端之间协作计算的通用方法。提供了 Controller 编程 API，可以使用它来创建协调客户端进行协作的工作流。FedAvg 是其中一个工作流的例子，另一个例子是循环权重传输。协作的核心概念是“任务”的概念。FL 控制器将任务（例如带有模型权重的深度学习训练）分配给一个或多个 FL 客户端，并处理客户端返回的结果（例如模型权重更新）。控制器可以基于处理的结果和其他因素（例如预先配置的训练轮次数）向客户端分配额外的任务。基于任务的交互将持续进行，直到目标实现为止。

API 支持典型的控制器-客户端交互模式，如上图所示：NVFlare 作业执行。控制器是一个 Python 对象，负责控制或协调 Worker 来完成工作。控制器在 FL 服务器上运行，Worker 则能够执行任务，并在 FL 客户端上运行。任务可以以广播方式发送给多个客户端，发送给一个或多个指定客户端，或者以顺序方式中继给多个客户端。每种交互模式都有两种类型：等待模式（阻塞直到接收到客户端结果）或非等待模式。可以使用这些交互模式来创建创新的工作流。例如，ScatterAndGather 控制器（通常用于类似 FedAvg 的算法）使用了 broadcast_and_wait 模式，而 CyclicController 使用了 relay_and_wait 模式。

控制器 API 使得无需处理底层的通信问题。每个 FL 客户端充当一个 Worker，简单地执行分配给它的任务（例如模型训练），并将执行结果返回给控制器。在每个任务交互中，可以有可选的过滤器在将任务数据或结果传递给控制器（在服务器端）或任务执行器（客户端）之前对其进行处理。该过滤机制可以用于数据隐私保护（例如同态加密/解密或差分隐私），而无需修改训练算法。

在生产模式下，通过使用 Notebooks 或 NVFlare 控制台提交作业来进行联邦学习。通过命令，可以启动或停止特定客户端或整个系统、提交新作业、检查作业状态、通过克隆现有作业创建新作业等。克隆作业的功能还是比较刚需，比如配置了复杂的联邦学习流程，当你需要修改一个小点，不需要重新连接，只需克隆一个新的作业做微调就可以。

2.5 系统概念

NVFlare 系统是一个典型的客户端-服务器通信系统，包含一个或多个 FL 服务器、一个或多个 FL 客户端，以及一个或多个管理员客户端。FL 服务器开启两个端口用于与 FL 客户端和管理员客户端通信。FL 客户端和管理员客户端通过已开放的端口进行连接。FL 客户端和管理员客户端无需开放任何端口，也不直接相互通信。

以下是框架中的关键概念和对象概述，以及它们之间可以传递的信息。

Workers 和 Controller
        NVFlare 的协同计算是通过 Controller 和 Worker 的交互来实现的。

Shareable
        代表服务器和客户端之间通信的对象。技术上，Shareable 实现为一个 Python 字典，包含不同的信息，例如模型权重。

数据交换对象 (DXO)
        标准化通信双方之间传递的数据。可以将 Shareable 比作信封，DXO 比作信件。两者一起构成了要在通信双方之间共享的消息。

FLComponent
        所有 FL 组件的基类。执行器、控制器、过滤器、聚合器及其子类型都是 FLComponent。FLComponent 带有一些内置方法用于日志记录、事件处理、审计和错误处理。

Executors
        FL 客户端上的 FLComponent 类型，具有一个 execute 方法，它从输入的 Shareable 生成一个新的 Shareable。提供了单进程和多进程执行器，以实现不同的计算工作负载。

FLContext
        允许在 FL 组件之间传递数据。FLContext 在所有常见 FLComponent 类型的每个方法中都可以使用。通过 FLContext，可以访问基础设施提供的服务，并与 FL 系统的其他组件共享数据。

通信驱动程序
        通信层抽象出来，因此可以为不同的部署场景实现可定制的通信驱动程序。默认情况下，使用 GRPC 进行基于任务的通信。然而，通信驱动程序可以替换为运行其他通信协议，例如 TCP。通信灵活性允许服务器中心化和点对点通信模式。这使得可以使用如 FedAvg 这样的广播与聚合工作流、分散式训练模式（如群学习），或点对点通信模式（如分割学习）。

过滤器
这是一种 FLComponent，具有 process 方法，用于在通信双方之间转换 Shareable 对象。过滤器可以在发送或接收 Shareable 数据之前进行额外处理。过滤器可以转换数据格式等，并且是实现数据隐私保护的主要机制：

ExcludeVars 用于排除 Shareable 中的变量。
PercentilePrivacy 用于通过百分位数截断权重。
SVTPrivacy 通过稀疏向量技术实现差分隐私。
同态加密过滤器用于安全聚合。

事件机制
框架带有事件机制，允许动态通知发送到所有事件处理器。该机制支持解耦组件之间基于数据的通信：当某个条件发生时，一个组件触发事件，其他组件可以监听该事件并处理事件数据。每个 FLComponent 自动是一个事件处理器。要监听和处理事件，只需实现 handle_event() 方法并处理所需的事件类型。事件代表系统逻辑执行过程中某些重要的时刻。例如，聚合前后或当重要数据变得可用时，例如选择了新的“最佳”模型。

以上所述框架概念或对象都是核心要素，这些要素基本构成了一个完整联邦计算框架。

3. 基于联邦计算框架实现联邦XGBOOST

3.1 联邦机器学习与联邦深度学习区别

在介绍联邦XGBOOST之前，我们先分析下联邦机器学习与联邦深度学习的区别。联邦传统机器学习和联邦深度学习之间的一个主要区别在于，对于传统的机器学习方法，"联邦"与"分布式"，甚至"集成"之间的界限可能比深度学习更加模糊。由于给定算法的特性及其API设计，这些概念有时是等同的。以XGBoost和SVM为例：在算法上，XGBoost可以将训练样本分配给多个计算节点，并基于每个节点收集的直方图构建树。这种过程可以直接应用于联邦学习场景，因为其通信成本是可接受的。在这种情况下，“联邦”相当于“分布式”学习。

在API方面，有些算法可能受到其实现的限制。以scikit-learn库中的SVM为例，虽然理论上SVM可以被形式化为一个迭代优化过程，但其API仅支持一次性“拟合”，并不支持单独调用优化步骤。因此，使用scikit-learn库实现的联邦SVM算法只能作为一个两步的过程。在这种情况下，“联邦”相当于“集成”学习。

为了说明这一点，提供基于树的联邦XGBoost（说明：这只是一种实现方式）的完整公式，并用下图进行了说明：

从定义上讲，XGBoost是一个顺序优化过程：每一步都向模型中添加一棵额外的树，以减少残差误差。因此，联邦XGBoost可以如下表述：每一轮联邦学习对应于本地层面的一个提升步骤。客户端在本地提升结束时将基于本地数据训练的新树共享给服务器。
模型表示为决策树或回归树。为了从所有客户端中聚合信息，服务器会将收到的所有树打包成一个“森林”，并将其添加到全局提升模型中。
在服务器更新全局模型后，每个客户端将继续提升过程，从全局提升森林模型开始学习一棵新的树。

3.2 安全联邦XGBoost设计

接下来我们开始分析下在nvflare中，联邦xgboost的实现方式。

3.2.1 横向安全模式

在横向XGBoost中，每个参与方拥有“平等地位”，即它们拥有部分样本的完整特征和标签，而联邦服务器只负责聚合操作，本身不持有任何数据。因此，从模型训练的角度来看，联邦服务器是一个“次要贡献者”，客户端则担心会将信息泄露给服务器。在这种设置下，保护措施主要针对本地直方图防止泄露给联邦服务器。

为了保护横向协作中的本地直方图，本地直方图在发送到联邦服务器进行聚合之前会被加密。聚合操作会在密文上进行，联邦服务器返回加密的全局直方图给客户端，客户端解密后再用于构建树。

3.2.2 纵向安全模式

在纵向XGBoost中，主动方持有标签信息，而被动方无法访问这些标签信息。标签被认为是整个过程中最有价值的资产。因此，从模型训练的角度来看，主动方是“主要贡献者”，它担心这些标签信息可能会泄露给被动客户端。在这种情况下，安全保护措施主要是防止标签信息泄露给被动客户端。

为了保护纵向协作中的标签信息，每轮XGBoost后，主动方会在计算每个样本的梯度后对其进行加密，然后将密文发送给被动方。被动方收到加密的梯度后，根据各自的特征分布对其进行累积。生成的累积直方图会返回给主动方，主动方解密后进一步用于构建决策树。

3.2.3 解耦加密与处理器接口

在设计上，XGBoost的通信通过通信层通过本地gRPC处理程序进行路由。从通信的角度来看，以前XGBoost内部的直接消息现在由FL通信器处理——它们通过FL系统与XGBoost之间的通信变成了“外部通信”。这为在XGBoost内部（在进入FL通信器之前）和FL系统内（通过FL通信器）执行消息操作提供了灵活性。XGBoost插件【4】将使用C++实现，而FL系统通信器将使用Python实现。设计并开发一个处理器接口，以便通过对特定同态加密方法和协作模式进行实现，正确连接这两者。

在接收到来自XGBoost的特定MPI调用后，各对应方调用接口进行数据处理（序列化等），提供必要的信息：g/h对或本地G/H直方图。处理器接口执行必要的处理（和加密），并将结果作为处理后的缓冲区返回。每个方随后将消息转发到FL系统侧的本地gRPC处理程序。经过涉及消息路由和计算的FL通信后，每个方在MPI调用时接收结果缓冲区。每个FL方将接收到的缓冲区发送到处理器接口进行解释。接口执行必要的处理（反序列化等），恢复正确的信息，并将结果发送回XGBoost以进行进一步计算。

3.2.4 系统设计

通过安全解决方案、通信模式和处理器接口，安全联邦XGBoost的示例设计如下，包括纵向和横向两种模式。

3.2.4.1 纵向联邦

这里需要注意下，active party是标签持有方(等同于fate中的guest方)，passive party是特征数据提供方（等同于fate中的host方）。

主动方首先计算其拥有的标签信息的g/h值。
g/h数据将发送到处理器接口，使用基于C++的加密工具库进行加密（这里涉及到使用GPU来加速密文处理），并通过FL通信发送到被动方。
被动方根据本地特征分布提供直方图计算所需的索引信息，处理器接口将对接收到的E(g/h)进行聚合（这里涉及到使用GPU来加速密文处理）。
结果E(G/H)将通过FL消息路由发送到主动方。
在主动方侧由处理器接口解密后，可以使用全局直方图信息进行树构建。

3.2.4.2 横向联邦

所有参与方通过处理器接口将其本地G/H直方图发送到FL端，在这个设计中，处理器接口仅执行缓冲区准备，而不涉及任何复杂的处理步骤。
在发送到联邦服务器之前，G/H直方图将在本地gRPC处理器中使用基于Python的加密工具库进行加密（这里涉及到使用GPU来加速密文处理）。
联邦服务器将对接收到的部分E(G/H)进行安全聚合（这里涉及到使用GPU来加速密文处理），并将全局E(G/H)分发给每个客户端，在那里全局直方图将被解密，并用于进一步的树构建。