在大型语言模型时代，通过消息量化和流式传输实现高效的联邦学习

联邦学习 (Federated Learning, FL) 已成为一种在分布式数据源中训练机器学习模型的有前景的方法，同时还能保护数据隐私。但是，在平衡模型要求和通信能力时，FL 面临着与通信开销和本地资源限制相关的重大挑战。

特别是在当前的大语言模型 (LLMs) 时代，FL 在部署具有数十亿参数的 LLMs 时面临着计算挑战。这些模型的庞大规模加剧了通信和内存限制。由于带宽限制，一次性传输完整的模型更新可能不可行，并且本地内存限制可能会使处理大型模型进行通信具有挑战性。解决这些问题需要创新策略。

NVIDIA FLARE 是一款与领域无关、开源且可扩展的联邦学习 SDK，通过引入可靠的通信功能、对多个并发训练作业的支持以及针对可能因网络条件而中断的作业的鲁棒性，增强了现实世界的联邦学习体验。

NVFlare 2.4.0 版本引入了流式传输 API，有助于传输超过 gRPC 规定的 2GB 大小限制的对象。它添加了一个新的流层，旨在可靠地处理大数据消息的传输。

借助流式传输 API，您将不再受到 gRPC 2-GB 大小限制的限制。然而，随着先进模型的日益壮大，两个挑战正在成为使用 LLM 的 FL 工作流的瓶颈：

• 默认 fp32 精度下的传输消息大小
• 分配本地内存，用于在传输过程中固定 object

为了实现更高效、更可靠的联邦管道，我们在 NVFlare 2.6.0 中引入了两项关键技术，以促进消息大小缩减和内存高效传输：

• 消息量化： 使用 NVFlare 过滤器 实现量化和去量化，并将其添加到联合方案中，从而减少传输过程中的消息大小。
• 容器和文件串流： 串流功能基于 ObjectStreamer 实现。我们支持两种对象类型 (容器和文件) ，并开发了 ObjectRetriever 类，以便更轻松地与现有代码集成。

流式传输功能：减少本地内存占用

FL 的主要瓶颈之一是远程参与者和服务器之间的模型更新交换。这些消息的大小可能非常大，令人望而却步，导致延迟和带宽消耗增加。鉴于最近的 LLMs 训练精度降低，NumPy 格式下的默认 fp32 消息精度甚至可能会人为地增大消息大小。

在本例中，我们实现了两项功能：native tensor transfer 和消息量化，通过实现原生训练精度，以及降低传输更新的精度和压缩消息大小，提供高效的消息传递解决方案。

图 1 展示了使用此 过滤器机制 实现量化和去量化的情况。在传输前对输出模型权重执行量化，而去量化在另一端接收消息时恢复原始精度。

这种实施有两个好处：

• 用户无需更改代码。相同的训练脚本可通过简单的配置设置用于和不用于消息量化
• 训练和聚合均以原始精度（而非量化数据）执行，从而最大限度地减少消息量化对训练过程的潜在影响。

我们使用直接裁剪和投射将 fp32 转换为 fp16，并利用 bitsandbytes 执行 8– 和 4-bit 量化。借助这些新功能，我们既支持 NumPy 数组 (之前的默认值) ，也支持直接用于训练 LLM 的 PyTorch Tensors。

表 1 显示了不同精度下 1B 参数 LLM 的消息大小 (以 MB 为单位) 。有关训练损失曲线对齐的更多信息，请参阅通过 Hugging Face SFT/PEFT API 进行 LLM 调整的示例。

通过应用消息量化技术，FL 可以显著节省带宽，并在我们的实验中使用 Supervised Fine-Tuning (SFT) 训练 LLM。

如图 2 所示，message quantization 不会在训练损失方面牺牲模型收敛质量。

流式传输功能：减少本地内存占用

FL 的另一个关键挑战是用于发送和接收消息的内存开销。

在默认设置下，要发送模型，您需要额外的内存来准备和接收模型块，这需要加倍本地内存使用量。必须分配额外的内存来保存整个消息以重新组合对象，尽管传输本身是通过 1M 流式传输块完成的。

如果系统性能出色且模型大小适中，这种内存开销是可以承受的，但如果考虑使用 70B 或更大的参数模型，则可能会迅速耗尽可用的系统内存。 70B 模型 的大小可以为 140 GB。要加载并发送它，您需要 140 + 140 = 280 GB 的内存。

尽管整个 LLM 参数字典可能非常庞大，但在分解为单个层和项目时，每层的最大大小要小得多，通常在 1 GB 左右。升级的流式传输功能通过两项新功能解决了内存使用方面的挑战：

• 对象容器流： 以增量方式处理和传输模型，而无需一次性将整个对象存储在内存中。容器流式传输一次对一个对象的一个项目 (例如包含模型权重的字典) 进行序列化。在之前的示例中，如果 ContainerStreamer 整体发送最大容量为 1 GB 的 140 GB 模型，则为 280 GB，而加载和发送 ContainerStreamer 只需要 140 + 1 = 141 GB 的内存。
• 文件流：流式传输文件而非结构化对象容器。文件流式传输逐块读取文件，并且仅消耗容纳一个数据块所需的内存。FileStreamer 所需的额外内存与模型大小或最大项目大小无关，并且仅依赖于文件 I/O 设置，这可以将传输内存占用率降至最低并实现无限流式传输。在这种情况下，无需加载模型，因此您可以根据需要选择进一步节省内存使用量。

在图 3 中，绿色框显示了必须为消息传输分配的最大本地内存。如图所示，常规传输必须为整个模型分配内存，因此随着模型的扩大，它可以是无限的。

对于对象容器，内存的大小与最大层相同，而最大层通常以第一层和最后一层为界。对于文件，内存需求独立于模型结构，并且可针对任何文件进行配置。

通过在 FL 中调整流式传输，您可以将更新分解为更小的块并按顺序进行处理，从而提高内存效率。流式传输可减少峰值内存占用，在优化计算资源的同时实现 FL。

借助此解决方案，您甚至可以实现实时处理，使设备能够在继续计算的同时传输部分更新，从而提高响应速度并减少空闲时间。在接收端，更新策略还可以从自适应传输中受益，其中可以根据网络条件和客户端可用性以不同粒度发送更新。

表 2 显示了一次性发送 1B 模型与本地模拟的内存比较。我们记录了系统内存占用情况，并比较了三种设置（regular、container streaming 和 file streaming）的峰值内存使用情况。

您可以看到，通过使用 streaming，memory usage 显著减少，尤其是 file streaming。但是，由于 file I/O 效率问题，file streaming 可能需要较长时间才能完成工作。

流增强功能尚未集成到高级 API 或现有的 FL 算法控制器和执行程序中。但是，您可以按照此 Streaming 示例构建自定义控制器或执行程序，以利用此功能。

总结

在本文中，我们展示了如何通过将消息量化和流式传输功能集成到 FL 框架来缓解通信瓶颈和内存限制。通过升级功能，我们提高了 Federated Learning 的效率和可扩展性。随着这些技术的不断发展，它们将在支持在不同环境中实际部署 FL 方面发挥至关重要的作用。

有关更多信息，请参阅以下资源：

• 在 GitHub 上的 / NVFlare 教程
• GitHub 上的/NVIDIA/NVFlare 量化示例
• GitHub 上的 /NVIDIA/NVFlare 流式传输示例
• NVIDIA FLARE 开发者门户
• 医学影像领域的 Federated Learning：增强数据隐私并推进医疗健康 GTC 2025 会议

要联系 NVIDIA FLARE 团队，请联系 federatedlearning@nvidia.com 。