跨 Hugging Face Jobs 的 LoRA 异步 GRPO

Hugging Face 已在 AsyncGRPOTrainer（TRL v1.14）中实现了 LoRA 支持，允许训练器仅将小型 LoRA 适配器同步到 vLLM，而非完整的模型权重。通过利用 Hugging Face 存储桶（Storage Buckets）作为共享文件系统，并配合用于请求路由的自定义代理，该架构使得训练和推理可以在不同的机器（Hugging Face Jobs）上运行，而无需 NCCL 或共享本地磁盘。

架构：通过存储桶进行分布式同步

AsyncGRPOTrainer 现在支持仅同步适配器的路径，从而消除了将张量直接发送到 vLLM 的需求。相反，训练器将 LoRA 适配器保存到存储桶中的特定目录，执行原子重命名，并通过 /v1/load_lora_adapter 端点通知 vLLM。

由于 Hugging Face Jobs 可以使用 hf-mount 将存储桶挂载为 FUSE 文件系统，因此训练器和 vLLM 副本可以在不同的虚拟机上共享相同的绝对路径。这消除了训练器和推理服务器共享物理节点或密集集群网络的需求。

任务布局

该系统由三个主要组件组成：

训练器任务（Trainer Job）：使用 LoRA 和 FSDP 运行 AsyncGRPOTrainer。
vLLM 任务（vLLM Jobs）：多个副本提供基础模型服务，并从存储桶加载最新的适配器。
代理服务器（Proxy Server）：一个运行在训练器任务上的小型基于 asyncio 的代理，负责处理身份验证头并将请求路由到 vLLM 副本。

代理：KV 前缀路由与广播

为了最大化多个 vLLM 副本之间的效率，使用自定义代理来管理请求分发和状态同步。

基于 KV 前缀的路由

为了避免冗余的预填充（prefill）计算，代理根据 KV 缓存前缀路由请求。它将提示词拆分为 16 个 token 的块，并计算以适配器名称为种子的链式哈希。

路由器跟踪哪个副本已处理了哪个块哈希。如果请求的提示词匹配特定副本上已缓存的前缀（且该副本未过载），则请求将被路由到该副本（即“亲和性命中”）。这防止了系统在多次 rollout 中对相同提示词重新计算预填充，这对于 GRPO 至关重要，因为 GRPO 会为单个提示词生成多个补全。

状态广播

由于每个 vLLM 副本都是一个独立的任务，代理通过向所有副本广播状态变更请求（如适配器加载、暂停和恢复）来确保一致性。这确保了特定的策略版本名称在整个集群中引用相同的权重。

性能优化与瓶颈分析

使用 sail/Sanity-Test-R1D-1.5B 数据集和 Qwen/Qwen2.5-Math-1.5B 模型，Hugging Face 进行了五次实验运行以优化流水线。结果表明，异步强化学习（Async RL）中的瓶颈可能在训练和生成之间转移。

关键优化

Token 预算批处理（Token-Budget Batching）：从每设备训练批大小 1 转向 token 预算批处理（例如 token_budget=16384），通过将多个序列打包到每一行中，减少了微批次的数量，从而将 MFU 从 3.9% 提高到 19%。
禁用梯度检查点（Disabling Gradient Checkpointing）：对于较小的模型（1.5B），禁用梯度检查点通过消除冗余的前向传播减少了前向+反向传播时间，将瓶颈从训练器转移到了生成副本。
增加在途请求（Increasing In-Flight Requests）：提高 max_inflight_tasks（例如设为 384）使系统能够充分利用多个 vLLM 副本，防止客户端并发限制降低吞吐量。

最终结果

通过结合这些优化，完成 500 步的总时间从 3 小时 27 分钟减少到 53 分钟（加速 3.9 倍）。

指标	运行 1（基线）	运行 5（优化后）
实际耗时（Wall Clock Time）	3 小时 27 分钟	53 分钟
中位步长耗时（Median Step Time）	22.9 秒	4.8 秒
MFU（前向/反向）	3.9%	23.5%
训练样本数（Samples Trained）	64,000	84,078
平均陈旧度（Mean Staleness）	1.5 个版本	2.0 个版本

技术实现细节

vLLM 版本：固定为 v0.27.1 以兼容运行时 LoRA 端点。
适配器插槽（Adapter Slots）：为了支持 max_staleness=4，vLLM 配置为 --max-loras 6，以确保在切换期间当前策略和之前的版本保持加载状态。
一致性（Consistency）：使用版本化的适配器名称，以防止 KV 缓存错误地匹配由旧策略版本生成的前缀。