NVIDIA HierarchicalKV(测试版)

关于 HierarchicalKV

HierarchicalKV 是 NVIDIA Merlin 的一部分，提供层次化的键值存储以满足推荐系统的需求。

HierarchicalKV 的主要功能是将键值对（特征-嵌入）存储在 GPU 的高带宽内存（HBM）和主机内存中。

您也可以将该库用于通用键值存储。

优势

在构建大型推荐系统时，机器学习（ML）工程师面临以下挑战：

需要使用 GPU，但单个 GPU 的 HBM 对于扩展到数 TB 的大型深度学习推荐模型（DLRM）来说太小。
在越来越大的 CPU 集群中改善通信性能变得越来越困难。
难以通过自定义策略高效控制有限 HBM 的消耗增长。
大多数通用键值库提供的 HBM 和主机内存利用率较低。

HierarchicalKV 缓解了这些挑战，并通过以下优势帮助推荐系统中的机器学习工程师：

支持同时在 HBM 和主机内存上训练大型推荐系统模型。
通过完全绕过 CPU 并减少通信工作负载来提供更好的性能。
实现基于 LRU 或自定义策略的表大小限制策略。这些策略通过 CUDA 内核实现。
在接近 1.0 的高工作状态负载因子下运行。

关键思想

桶是局部有序的
分别存储键和值
将所有键存储在 HBM 中
内置和可自定义的淘汰策略

HierarchicalKV 使 NVIDIA GPU 更适合训练搜索、推荐和广告的大型和超大型模型。该库简化了构建、评估和提供复杂推荐模型的常见挑战。

API 文档

以下是主要的类和结构体，但建议阅读源代码中的注释：

常规 API 文档请参考 API 文档

API 成熟度矩阵

"行业验证"表示该 API 已在至少一个真实场景中经过充分测试和验证。

名称	描述	功能
insert_or_assign	插入或更新指定键。<br>当桶满时，用最小分数覆盖一个键。	行业验证
insert_and_evict	插入新键，当桶满时淘汰最小分数的键。	行业验证
find_or_insert	搜索指定键，如果未找到则插入。	充分测试
assign	更新每个键，如果未找到则跳过。	充分测试
accum_or_assign	搜索并更新每个键。如果找到，将值作为增量添加到原始值。<br>如果未找到，则直接更新。	充分测试
find_or_insert*	搜索指定键并返回值的指针。如果未找到，先插入。	充分测试
find	搜索指定键。	行业验证
find*	搜索并返回值的指针，线程不安全但性能高。	充分测试
export_batch	导出一定数量的键-值-分数元组。	行业验证
export_batch_if	导出符合特定条件的一定数量的键-值-分数元组。	行业验证
warmup	将热门键值从主机内存移动到 HBM	2023年6月15日

淘汰策略

引入"分数"来定义每个键的重要性，分数越大越重要，被淘汰的可能性越小。只有当桶满时才会发生淘汰。 score_type 必须是 uint64_t。更多详情，请参考 EvictStrategy 类。

名称	`分数`定义
Lru	设备时钟，以纳秒为单位，可能与主机时钟略有不同。
Lfu	频率增量由调用者通过"插入类"API 的 `scores` 输入参数提供，作为频率的增量。
EpochLru	高 32 位是通过 `global_epoch` 输入参数提供的全局纪元，<br>低 32 位等于 `(device_clock >> 20) & 0xffffffff`，粒度接近 1 毫秒。
EpochLfu	高 32 位是通过 `global_epoch` 输入参数提供的全局纪元，<br>低 32 位是频率，<br>频率在达到 `0xffffffff` 的最大值后将保持不变。
Customized	完全由调用者通过"插入类"API 的 `scores` 输入参数提供。

注意：
- "插入类"API 指 insert_or_assign、insert_and_evict、find_or_insert、accum_or_assign 和 find_or_insert 等 API。
- global_epoch 应由调用者维护，并作为"插入类"API 的输入参数。

配置选项

建议保持以 * 结尾的选项的默认配置。

名称	类型	默认值	描述
init_capacity	size_t	0	哈希表的初始容量。
max_capacity	size_t	0	哈希表的最大容量。
max_hbm_for_vectors	size_t	0	向量的最大 HBM，以字节为单位。
dim	size_t	64	值向量的维度。
max_bucket_size*	size_t	128	每个桶的长度。
max_load_factor*	float	0.5f	重新哈希前的最大负载因子。
block_size*	int	128	CUDA 内核的默认块大小。
io_block_size*	int	1024	IO CUDA 内核的块大小。
device_id*	int	-1	设备 ID。设置为 `-1` 时内部管理。
io_by_cpu*	bool	false	指示是否由 CPU 处理 IO 的标志。
reserved_key_start_bit	int	0	64 位中保留键的起始位偏移。

更多详情请参考 HashTableOptions 结构体。

保留键值

默认情况下，键值 0xFFFFFFFFFFFFFFFD、0xFFFFFFFFFFFFFFFE 和 0xFFFFFFFFFFFFFFFF 被保留用于内部使用。如果您想使用上述键值，请修改 options.reserved_key_start_bit。 reserved_key_start_bit 的有效范围是 0 到 62。默认值为 0，对应上述默认保留键值。当 reserved_key_start_bit 设置为非 0 值时，任何保留键值的最低有效位（第 0 位）始终为 0。
设置 reserved_key_start_bit = 1:
- 此设置为保留键值保留最低两位 1 和 2。
- 二进制表示中，最后四位范围从 1000 到 1110。这里，最低有效位（第 0 位）始终为 0，第 3 到 63 位设置为 1。
- 新的保留键值的十六进制表示如下：
  - 0xFFFFFFFFFFFFFFFE
  - 0xFFFFFFFFFFFFFFFC
  - 0xFFFFFFFFFFFFFFF8
  - 0xFFFFFFFFFFFFFFFA
设置 reserved_key_start_bit = 2:
- 此配置将第 2 和第 3 位保留为保留键值。
- 最后五位的二进制表示范围从 10010 到 11110，最低有效位（第 0 位）始终为 0，第 4 到 63 位设置为 1。
如果您更改了 reserved_key_start_bit，在保存/加载时应使用相同的值更多详情，请参考 init_reserved_keys

使用方法：

#include "merlin_hashtable.cuh"


using TableOptions = nv::merlin::HashTableOptions;
using EvictStrategy = nv::merlin::EvictStrategy;

int main(int argc, char *argv[])
{
  using K = uint64_t;
  using V = float;
  using S = uint64_t;
  
  // 1. 定义表并使用 LRU 淘汰策略。
  using HKVTable = nv::merlin::HashTable<K, V, S, EvictStrategy::kLru>;
  std::unique_ptr<HKVTable> table = std::make_unique<HKVTable>();
  
  // 2. 定义配置选项。
  TableOptions options;
  options.init_capacity = 16 * 1024 * 1024;
  options.max_capacity = options.init_capacity;
  options.dim = 16;
  options.max_hbm_for_vectors = nv::merlin::GB(16);
  
  
  // 3. 初始化表的内存资源。
  table->init(options);
  
  // 4. 使用表进行操作。
  
  return 0;
}

使用限制

key_type 必须是 int64_t 或 uint64_t。
score_type 必须是 uint64_t。

贡献者

HierarchicalKV 由 NVIDIA Merlin 团队和 NVIDIA 产品最终用户共同维护，同时也欢迎公众贡献、错误修复和文档完善。 [贡献指南]

如何构建

基本上，HierarchicalKV 是一个仅包含头文件的库，以下命令仅用于创建基准测试和单元测试的二进制文件。

您的环境必须满足以下要求：

CUDA 版本 >= 11.2
NVIDIA GPU，计算能力为 8.0、8.6、8.7 或 9.0
GCC 支持 C++17 标准或更高版本
Bazel 版本 >= 3.7.2（仅用于 Bazel 编译）

使用 cmake

git clone --recursive https://github.com/NVIDIA-Merlin/HierarchicalKV.git
cd HierarchicalKV && mkdir -p build && cd build
cmake -DCMAKE_BUILD_TYPE=Release -Dsm=80 .. && make -j

调试版本：

cmake -DCMAKE_BUILD_TYPE=Debug -Dsm=80 .. && make -j

基准测试：

./merlin_hashtable_benchmark

单元测试：

./merlin_hashtable_test

使用 bazel

请勿在 git clone 时使用 --recursive 选项。
如果使用自定义 Docker 镜像，请预先修改 .bazelrc 文件中的环境变量。
强烈推荐使用 nvcr.io/nvidia/tensorflow 上维护的 Docker 镜像。

拉取 Docker 镜像：

docker pull nvcr.io/nvidia/tensorflow:22.09-tf2-py3
docker run --gpus all -it --rm nvcr.io/nvidia/tensorflow:22.09-tf2-py3

在 Docker 容器中编译：

git clone https://github.com/NVIDIA-Merlin/HierarchicalKV.git
cd HierarchicalKV && bash bazel_build.sh

基准测试：

./benchmark_util

基准测试和性能（进行中）

GPU: 1 x NVIDIA A100 80GB PCIe: 8.0
键类型 = uint64_t
值类型 = float32 * {dim}
每次操作的键值对数量 = 1048576
淘汰策略：LRU
λ：负载因子
find* 表示直接返回值地址的 find API。
find_or_insert* 表示直接返回值地址的 find_or_insert API。
吞吐量单位：十亿键值对/秒

纯 HBM 模式：

dim = 8，容量 = 1.28 亿键值对，HBM = 4 GB，HMEM = 0 GB

λ	insert_or_assign	find	find_or_insert	assign	find*	find_or_insert*	insert_and_evict
0.50	1.093	2.470	1.478	1.770	3.726	1.447	1.075
0.75	1.045	2.452	1.335	1.807	3.374	1.309	1.013
1.00	0.655	2.481	0.612	1.815	1.865	0.619	0.511

λ	export_batch	export_batch_if	contains
0.50	2.087	12.258	3.121
0.75	2.045	12.447	3.094
1.00	1.950	2.657	3.096

dim = 32，容量 = 1.28 亿键值对，HBM = 16 GB，HMEM = 0 GB

λ	insert_or_assign	find	find_or_insert	assign	find*	find_or_insert*	insert_and_evict
0.50	0.961	2.272	1.278	1.706	3.718	1.435	0.931
0.75	0.930	2.238	1.177	1.693	3.369	1.316	0.866
1.00	0.646	2.321	0.572	1.783	1.873	0.618	0.469
λ	export_batch	export_batch_if	contains
-----:	-------------:	----------------:	---------:
0.50	0.692	10.784	3.100
0.75	0.569	10.240	3.075
1.00	0.551	0.765	3.096