Hugging Face TRL 和 RapidFire AI 集成

Hugging Face TRL 已正式与 RapidFire AI 集成,以加速大型语言模型(LLM)的微调和后训练。此集成使开发者能够在单个 GPU 上甚至无需显著更改代码或增加 GPU 硬件要求的情况下,并发比较多个训练配置。

加速实验吞吐量

RapidFire AI 相比顺序配置测试,将实验速度提升了 16–24 倍。这种加速是通过自适应的基于块的调度和执行方案实现的,该方案允许多个 TRL 配置并发运行,使团队能够更快地达到最佳评估指标。

在 NVIDIA A100 40GB GPU 上使用 TinyLlama-1.1B 和 Llama-3.2-1B 模型进行的内部基准测试展示了达到最佳训练损耗时间的以下加速:

场景 顺序时间 RapidFire AI 时间 加速倍数
4 个配置, 1 GPU 120 min 7.5 min 16×
8 个配置, 1 GPU 240 min 12 min 20×
4 个配置, 2 GPU 60 min 4 min 15×

技术架构与核心能力

RapidFire AI 通过几个关键的技术机制优化 GPU 利用率和开发者工作流程:

自适应基于块的并发训练

RapidFire AI 将数据集划分为指定数量的块,并在块边界处通过 GPU 循环不同的 LLM 配置。此机制相比顺序训练能够更早地在所有配置上提供评估指标的增量信息,从而促进更快的比较决策。

共享内存编排

为了保持稳定性和性能,系统使用基于高效共享内存的适配器和模型溢出/加载机制进行自动检查点。调度器利用这些共享内存机制在可用 GPU 上自动编排配置,从而消除了手动管道的需求。

交互式控制操作(IC Ops)

用户可以通过实时仪表板管理进行中的实验。IC Ops 允许开发者:

  • Stop:终止表现不佳的配置以节省资源。
  • Resume:重新启动暂停的运行。
  • Delete:删除不必要的运行。
  • Clone-Modify:基于有希望的运行创建新配置,修改超参数,并可选择从父权重进行热启动。

TRL 集成与实现

RapidFire AI 为 TRL 训练器提供即插即用的包装器,使用户在获得并发能力的同时保持其现有的 TRL 思维模型。支持的训练器包括:

  • SFT:使用 RFSFTConfig
  • DPO:使用 RFDPOConfig
  • GRPO:使用 RFGRPOConfig

实现示例

要在单个 GPU 上实现并发训练,用户可以使用 RFModelConfigRFLoraConfig 定义一个 config_set,然后通过指定的 num_chunks 使用 experiment.run_fit 执行实验。例如,在 2-GPU 设置中,RapidFire AI 可以将比较决策的时间从大约 15 分钟(顺序)减少到 5 分钟(并发),同时将 GPU 利用率从 60% 提高到超过 95%。

开始使用

RapidFire AI 可以通过 PyPI 使用 pip install rapidfireai 进行安装。初始设置需要通过 Hugging Face CLI 进行身份验证,并执行 rapidfireai initrapidfireai start。系统提供了一个基于 MLflow 的仪表板,可通过 http://localhost:3000 访问以进行实时监控和控制,未来计划增加对 Trackio、W&B 和 TensorBoard 的支持。

Sources