PPO和其他LLM微调技术需要在训练过程中使用自回归生成作为部分。当使用vLLM加速训练循环的自回归生成部分时,是否有一种有效的方式来更新LLM的权重?具体来说,在LoRA微调的情况下,是否有一种方法可以高效地交换适配器,而不必将其保存到文件系统?
高效的LoRA适配器更新
可能的解决方法无需进行任何代码更改:将适配器保存到内存文件系统中(例如,/dev/shm
),并在每个LoRARequest中指向该目录。这个解决方法的优点是:
- 避免了磁盘读/写瓶颈和SSD磨损。
- 仍然会产生safetensors序列化和反序列化的开销。
建议的更改:修改LoRARequest以允许将适配器指定为Tensor的字典。
- 修改LoRARequest类定义
- 将
lora_local_path: str
标记为可选 - 添加新的可选
lora_tensors: dict[str, torch.Tensor]
属性。 - 修改WorkerLoRAManager
_load_lora
实现(vllm/lora/worker_manager.py) - 验证给定的LoRARequest指定了
lora_local_path
和lora_tensors
中的一个。 - (可选)将检查
unexpected_modules
的逻辑移动到单独的方法中。 - 如果在LoRARequest中提供了
lora_tensors
: - 在给定的Tensor字典中检查
unexpected_modules
。 - 用
from_lora_tensors
替换from_local_checkpoint
。
替代方法:非LoRA参数更新
- OpenRLHF通过覆盖
hf_model_weights_iterator
并为字典中的每个Tensor调用load_weights
来用内存Tensor替换vLLM模型参数。(来源,补丁)
其他上下文
LLM微调目标(如PPO)在训练过程中需要自回归文本生成,要求使用相对较新的模型进行生成。
截至v0.4.0 vLLM发布版本,当示例化一个vLLM LoRARequest
时,通过lora_local_path: str
属性指定LoRA适配器。(来源)在上面的LoRA PPO示例中,如果vLLM示例与peft
训练循环在同一台机器上,将适配器权重的新副本发送给vLLM需要以下步骤:
- 调用
peft.PeftModel.save_pretrained
将适配器Tensor状态字典(作为folder_name/adapter_model.safetensors
)保存到磁盘上的本地路径。在此幕后,此方法将执行以下操作: - 调用
peft.utils.get_peft_model_state_dict
获取Tensor字典,然后 - 调用
safetensors.torch.save_file
将loraTensor字典序列化到文件系统。(序列化开销) - 示例化一个新的vLLM
LoRARequest
并将lora_local_path
属性设置为更新后的值。 - 将此
LoRARequest
发送给vLLM引擎。在此幕后,vLLM将执行以下操作: - 调用
LoRAModel.from_local_checkpoint
(来源) - 验证peft配置中列出的所有
target_modules
是否受支持。 - 从文件系统将loraTensor字典加载到CPU内存中。(反序列化开销)
- 如果提供了其他嵌入Tensor,则将它们加载到CPU内存中。
- 调用
LoRAModel.from_lora_tensors
(来源)以示例化LoRAModel。
如果采用的替代方案被采纳,新的工作流程如下:
- 在LoRA模型上调用
peft.utils.get_peft_model_state_dict
以获取loraTensor字典(与上述解决方法中写入磁盘的操作相同)。 - 示例化一个新的vLLM LoRARequest并包含指向该loraTensor字典的指针。
- 将此
LoRARequest
发送给vLLM引擎。在此幕后,vLLM将执行以下操作: - 调用
LoRAModel.from_lora_tensors
(来源)以示例化更新后的LoRAModel。
4条答案
按热度按时间k10s72fa1#
@jacobthebanana 很高兴知道Open RLHF做了这样的事情。你知道有没有关于权重广播的最小示例吗?
ki1q1bka2#
@jacobthebanana 很高兴知道Open RLHF做了这样的事情。你知道有没有一个最小化的权重广播示例吗?
我的一位同事在OpenRLHF仓库中分享了这个例子- examples/train_ppo_ray.py 。对于我最感兴趣的用例,这种方法的主要缺点是它需要大量的GPU。具体来说,在这个设置中,vLLM引擎需要在其自己的GPU集合上运行-与运行反向传播的GPU分开。
参考一下,OpenRLHF完整秩vLLM热插拔逻辑可以在 openrlhf/trainer/ray/vllm_worker_wrap.py 中找到,它在openrlhf/trainer/ray/ppo_actor.py中被调用。另一个挑战是,在同一组GPU上直接运行Torch FSDP和使用Ray的vLLM并不简单。当 pull request #3466 为vLLM合并时,这可能会变得更容易。我在工作中的团队基于该拉取请求中提出的更改构建了一个基于LoRA权重“广播”的概念证明-使用上面提到的
/dev/shm
解决方法。如果你对此感兴趣,我很乐意分享更多关于这项努力的信息。(另外,显然我的GitHub电子邮件通知没有正确设置。抱歉回复晚了。)
ppcbkaq53#
太酷了,谢谢你的回复!我觉得一个非常有影响力的项目是直接训练vLLM模型。
在在线RLHF方面,我还实现了将模型放置在特定设备vwxyzjn#1上,然后将其应用于TRL:huggingface/trl#1540的功能。这个想法是在第8个GPU上加载vLLM模型,并使用剩余的GPU进行训练。
bqf10yzr4#
这太酷了,谢谢你的回复!我认为一个非常有影响力的项目是直接训练vLLM模型。
在在线RLHF方面,我也使得将模型放置在特定设备vwxyzjn#1上并将其应用于TRL:huggingface/trl#1540成为可能。的想法是在第8个GPU上加载vLLM模型,并使用剩余的GPU进行训练。
这看起来是一种在训练过程中实现推理和热插拔的非常优雅的方法!的确-运行推理所需的GPU内存比训练少得多,而vLLM通过分页注意力进一步降低了内存要求。就运行vLLM引擎所需的内存而言,一个GPU应该足够了。
我必须承认我对HuggingFace
accelerate
并不特别熟悉。我看到你一直在调用model.load_weights-你知道这个过程是否需要经过CPU内存空间?我在想你是否观察到与权重传输相关的显著吞吐量限制。此外,你有关于在所有8个GPU上运行vLLM所需的工作的大致估计吗?我的团队一直在寻找方法,通过在与训练循环相同的设备上运行vLLM来充分利用我们的(有限数量的)GPU。因为torch FSDP需要对nccl的独占访问,我们最终不得不为vLLM和我们自己的训练逻辑创建多个 Package 器。
accelerate
(而不是FSDP)是否会成为一个更好的选择,使训练循环能够与vLLM逻辑并行运行?