FBTT-Embedding安装与使用指南

FBTT-Embedding安装与使用指南

FBTT-EmbeddingThis is a Tensor Train based compression library to compress sparse embedding tables used in large-scale machine learning models such as recommendation and natural language processing. We showed this library can reduce the total model size by up to 100x in Facebook’s open sourced DLRM model while achieving same model quality. Our implementation is faster than the state-of-the-art implementations. Existing the state-of-the-art library also decompresses the whole embedding tables on the fly therefore they do not provide memory reduction during runtime of the training. Our library decompresses only the requested rows therefore can provide 10,000 times memory footprint reduction per embedding table. The library also includes a software cache to store a portion of the entries in the table in decompressed format for faster lookup and process.项目地址:https://gitcode.com/gh_mirrors/fb/FBTT-Embedding

FBTT-Embedding 是一个基于张量列车(Tensor Train)的压缩库，专为大型机器学习模型（如推荐系统和自然语言处理）中的稀疏嵌入表设计。此库能显著减少模型大小，例如，在Facebook公开的DLRM模型中可达到100倍的压缩比率，且不影响模型性能。接下来，我们将详细解析项目的结构、启动文件以及配置相关的指导。

1. 项目目录结构及介绍

FBTT-Embedding的项目结构组织如下：

• src: 包含主要的源代码文件，实现Tensor Train的压缩逻辑和相关操作。

  • cpp: C++源代码，包括核心算法实现。
  • py: Python接口，使得在PyTorch环境中可以轻松调用这些功能。
  • cuda: 使用CUDA编写的GPU相关部分，优化计算性能。
  • ops: 操作相关，可能是封装了特定运算的模块。
  • test: 测试脚本，用于验证库的功能正确性。

• benchmark: 性能基准测试工具和脚本，如tt_embeddings_benchmark.py用于评估不同设置下的性能。

• setup.py: 安装脚本，用于将项目集成到Python环境中。

• LICENSE: 许可证文件，说明了项目的授权方式（MIT许可证）。

• README.md: 项目简介，包含了快速入门、限制和资源链接等重要信息。

2. 项目的启动文件介绍

• setup.py：这是安装FBTT-Embedding的关键脚本，通过执行python setup.py install命令，你可以将该库安装到你的Python环境。开发者可以在该文件内调整编译选项，以适配不同的硬件架构（如通过修改nvcc的-gencode参数支持Volta或Pascal架构）。

• run_tt_embeddings_benchmark.py：示例运行脚本，用于进行性能测试。通过这个脚本，你可以得到TT-Embeddings在特定配置下的前向传播和反向传播的时间、GFLOPS及带宽等性能指标。

3. 项目的配置文件介绍

FBTT-Embedding本身没有明确的传统配置文件（如.ini, .yaml等），但其配置和设置主要通过代码内的参数传递来完成。例如，在创建TTEmbeddingBag实例时，通过参数如num_embeddings, embedding_dim, tt_ranks, sparse, 和 optimizer等直接指定配置项。这意味着配置是动态的，依赖于初始化库函数时提供的参数。

为了自定义配置，用户需要直接在使用库的地方指定这些关键参数。例如，当需要改变压缩等级、优化器类型或是否启用软件缓存时，这些改动都应在实例化TTEmbeddingBag类时进行。

示例配置实例

在Python中配置并使用FBTT-Embedding可能如下所示：

from fbtt_embedding import TTEmbeddingBag
import torch

# 自定义配置
num_embeddings = 1000000
embedding_dim = 64
tt_p_shapes = [120, 90, 110]
tt_q_shapes = [4, 4, 4]
tt_ranks = [12, 14]
sparse = True
optimizer_type = 'SGD'
learning_rate = 0.05

# 创建嵌入层
embedding_layer = TTEmbeddingBag(
    num_embeddings=num_embeddings,
    embedding_dim=embedding_dim,
    tt_p_shapes=tt_p_shapes,
    tt_q_shapes=tt_q_shapes,
    tt_ranks=tt_ranks,
    sparse=sparse,
    optimizer=optimizer_type,
    learning_rate=learning_rate
)

# 假设的输入数据
input = torch.LongTensor([[0, 1], [2, 3]])
offsets = torch.LongTensor([0, 2])

# 使用层进行前向传播
output = embedding_layer(input, offsets)

请注意，实际应用中应根据具体需求调整上述参数，并确保系统环境已满足安装和运行所需的依赖，尤其是PyTorch和CUDA版本。

FBTT-EmbeddingThis is a Tensor Train based compression library to compress sparse embedding tables used in large-scale machine learning models such as recommendation and natural language processing. We showed this library can reduce the total model size by up to 100x in Facebook’s open sourced DLRM model while achieving same model quality. Our implementation is faster than the state-of-the-art implementations. Existing the state-of-the-art library also decompresses the whole embedding tables on the fly therefore they do not provide memory reduction during runtime of the training. Our library decompresses only the requested rows therefore can provide 10,000 times memory footprint reduction per embedding table. The library also includes a software cache to store a portion of the entries in the table in decompressed format for faster lookup and process.项目地址:https://gitcode.com/gh_mirrors/fb/FBTT-Embedding