参数服务器是推荐系统的重要组成部分,但是目前的训练端参数服务器由于高延迟和同步问题无法有效解决推理部署中模型过大的瓶颈。MerlinHugeCTR(以下简称HugeCTR)团队针对传统参数服务器的问题重新设计了一种分级推理端参数服务器,将GPU内存作为一级缓存,Redis集群作为二级缓存,RocksDB作为持久化层,极大提高了推理效率。HugeCTR团队将分多期为大家介绍此分级参数服务器的具体设计和细节,本期为系列的第一期。
引言
a.传统训练端参数服务器及其缺陷
传统参数服务器维护和同步模型参数仅用于训练,worker节点执行前向和后向计算。具体来说,在训练中:worker节点从server节点中拉取其相应的参数,进行前向计算,通过反向传播计算梯度,最后将这些梯度推送到服务器。在推理中,它只执行前两个步骤。如果部署在高性能设备集群中,worker节点的计算速度非常快,因此传统PS通常会遇到这两个瓶颈:
(1)server和worker之间的pull和push操作延迟;
(2)从worker节点收到梯度后,server节点中的参数同步问题。
由于GPU停顿、同步/一致性不足,GPU的计算结构很难通过使用基于CPU的实现的参数服务器来支持数据并行。在GPU内存中拟合完整模型以及小批量输入数据和中间网络状态的需要限制了可以训练和推理的模型的大小。同样的瓶颈也出现在推理部署中,因为推理节点也需要从集中的参数服务器组中拉取所有需要的模型参数。当请求包含节点未加载的参数时,节点需要再次同步从参数服务器拉取参数。当模型参数版本发生变化时,需要暂停推理服务,逐个节点更新参数。
使用基于CPU的参数服务器进行模型推理的部署时,上述问题非常明显,特别是推荐模型的部署。
b.HugeCTR推理端分级参数服务器
与其他系统不同,HPS进行了许多专门针对高效利用GPU的优化,包括分布式参数服务器分片,以实现GPU/CPU中许多庞大EmbeddingTable的并行推理,以及GPU友好的缓存、临时数据移动内存、内存管理机制。
图1HPS架构
c.推理端参数服务器的支持
●支持不同模型的混合部署:如DeepFM、DCN、DLRM、MMOE和序列模型(DIN、DIEN)
●支持推理的大输入数据量:Batch_size大于1K,look_upperrequest超过1000。
●支持更快的在线热部署:将完整模型更新/加载到推理节点进行服务只需不到10分钟(EmbeddingTable大小大于600G)
●支持资源隔离:在推理中隔离GPU的内存,确保推理服务在生产环境基于不同的隔离策略。例如一种支持单GPU/CPU的模型,通过第三方工具(如k8s)重启和隔离。
●支持不同模型独立的巨大EmbeddingTable:不同模型每个EmbeddingTable大小大于600G。
●支持单个节点的多级缓存:
●支持模型更新机制:
●支持容错和持久性:一个节点故障后参数和服务可以被恢复。
●支持在线学习:每分钟更新密集模型权重
HugeCTR分级参数服务器组件块
a.CPU分布式缓存
i.分布式Redis集群
Redis集群的同步查找:每个模型实例从本地化的GPU缓存中查找所需的Embeddingkey,这也会将丢失的Embeddingkey(在GPU缓存中找不到的key)存储到丢失的键缓冲区中。丢失的键缓冲区与Redis实例同步交换,Redis实例依次对任何丢失的Embeddingkey执行查找操作。因此,分布式Redis集群充当了二级缓存,可以完全替代本地化参数服务器来加载所有模型的完整Embeddingtable。
b.GPU缓存(EmbeddingCache)
i.异步/同步插入
我们支持将丢失的Embeddingkey异步插入到EmbeddingCache中。该功能可以通过配置文件中自定义的命中率阈值自动激活。当EmbeddingCache的真实命中率高于自定义阈值时,EmbeddingCache会异步插入缺失的key。反之则会以同步方式插入,以确保推理请求的高精度。通过异步插入方式,与之前的同步方式相比,在EmbeddingCache达到用户定义的阈值后,可以进一步提高EmbeddingCache的真实命中率。
ii.在线更新
我们支持将增量EmbeddingKey异步刷新到EmbeddingCache中。当稀疏模型文件需要更新到GPUEmbeddingCache时,会触发刷新操作。基于在线训练完成模型的模型版本迭代或增量参数更新后,需要将最新的Embeddingtable更新到推理服务器上的EmbeddingCache中。为了保证运行模型可以在线更新,我们将通过分布式事件流平台(Kafka)更新分布式数据库和持久化数据库。同时,GPUEmbeddingCache会刷新现有Embeddingkey的值,并替换为最新的增量Embeddingvector。
c.本地键值存储
i.本地RocksDB查询引擎
对于仍然无法完全加载到Redis集群中的超大规模Embddingtable我们将在每个节点上启用本地键值存储(RocksDB)。
RocksDB的同步查询:Redis集群客户端在分布式GPU缓存中查找Embeddingkey时,会记录丢失的Embeddingkey(在Redis集群中未找到的key),记录到丢失的keybuffer中。丢失的keybuffer与本地RocksDB客户端同步交换,然后将尝试在本地SSD中查找这些key。最终,SSD查询引擎将对所有模型缺失的Embeddingkey执行第三次查找操作。
对于已经存储在云端的模型存储库,RocksDB将作为本地SSD缓存,用于存储Redis集群无法加载的剩余部分。因此,在实践中,本地化的RocksDB实例充当了三级缓存。
HugeCTR分级参数服务器的配置和使用
a.训练端
i.配置模型名和Kafkabroker
在训练段,用户在CreateSolver时需要为当前训练的模型提供一个模型名,这个模型名将会被参数服务器用于区分不同模型。
同时,用户还需要配置Kafkabroker的端口和ip,用于将模型发送到到Kafka。
ii.增量模型导出接口
用户可以使用以上接口,将增量模型导出到Kafkabroker,参数服务器端将会自动消化Kafka的消息。
b.推理端
i.EmbeddingCache配置:
gpucache:用户可自由配置是否使用GPU缓存。
gpucacheper:用于决定Embeddingtable导入到GPU缓存的比例,默认为0.5。
hit_rate_threshold:用户自定义的阈值,将会决定GPU缓存的更新方式。
num_partitions:Embeddingtable将会被分为多个分片进行存储,这里用于指定分片数量。
overflow_policy:当缓存占满时,可选择随机移除或移除最老的Embedding。
overflow_margin:用于指定每个分片储存的最大Embedding数量。
overflow_resolution_target:用于指定每个分片移除Embedding的比例,取值为0到1之间。
initial_cache_rate:初始的缓存率。
需要配置服务器ip和端口,用户名以及密码,其他与hashmap/parallelhashmap相同。需要注意的是num_partitions必须大于等于redis节点的数量。
iv.RocksDB:
path:RocksDB存储数据的路径,由用户自行配置。
read_only:启用read_only后,RocksDB将无法更新,适用于静态Embedding的推理。
v.Kafka:
brokers:用于设置Kafka服务器的ip和端口。
max_receive_buffer_size:用于数据接收缓冲区大小,超过该大小将自动把参数更新送往数据存储层。
max_batch_size:用于设置每次批量数据发送的大小。
结语
在这里,要特别鸣谢腾讯PCG机器学习平台部推荐算法资深高级工程师骆兆楷(KarlLuo)对HugeCTR设计开发的大力支持与密切合作。
除此之外,NVIDIAMerlinHugeCTR团队正在积极招募C++以及CUDA工程师,欢迎各位有意向的同学积极申请!