合聚咖

合聚咖

LLM在IREE上进行多设备部署

admin

在IREE平台上部署LLMs面临多设备部署的挑战,包括优化数据流动、设备分配和内存限制。执行模型时,计算密集型任务通常倾向于GPU,而I/O密集型操作更适合CPU。GPU内存有限,因此大模型如LLM不能直接offload,需处理量化、磁盘管理和跨CPU-GPU的数据缓存。I/O瓶颈是关键问题,目标是通过流水线并行和异步执行来提高DAG的执行效率。

执行流程涉及流水线的多阶段划分,每个阶段分配到特定设备。目标在于减少设备间数据传输,保持各阶段计算时间均衡,同时最小化阶段间数据流动,并寻求最小化总执行时间。同构设备情况下,可采用多级算法对DAG进行分割。而在异构设备中,如CPU+GPU组合,需考虑不同特性的内存带宽和计算能力,确保数据能准确并高效地在设备间移动。

当内存充足时,可以分布式运行模型,无需频繁load/unload。而在内存有限的情况下,需要进行缓存策略,如静态缓存在编译时明确描述数据复制,而动态缓存则在运行时动态处理。权重压缩和量化是优化模型在不同设备上执行的手段,如GPTQ的量化效果已经证明了在小到桌面大到移动设备上的可行性。

为了支持多设备部署,IREE需要操作符来处理设备间的数据复制,同时考虑同步问题。磁盘可以作为设备模型,支持数据复制。在IREE中,通过MLIR操作描述数据在设备间的移动,现有操作如发送和接收仅限于同一运行时环境。统一内存的处理,如Grace-Hopper,需要考虑如何整合为一个逻辑设备。

最后,外部API的使用意味着模型可能被拆分为多个模块,数据传输在IREE运行时之外处理。这可能涉及在用户和IREE运行时之间增加一层驱动,以协调数据传输的调度。