PyTorch 2.2 中文官方教程（二十）（2）https://developer.aliyun.com/article/1482630
推荐系统
TorchRec 简介
原文：
pytorch.org/tutorials/intermediate/torchrec_tutorial.html译者：飞龙
提示
为了充分利用本教程，我们建议使用这个Colab 版本。这将使您能够尝试下面提供的信息。
请跟随下面的视频或在youtube上观看。
www.youtube.com/embed/cjgj41dvSeQ
在构建推荐系统时，我们经常希望用嵌入来表示产品或页面等实体。例如，参见 Meta AI 的深度学习推荐模型，或 DLRM。随着实体数量的增长，嵌入表的大小可能超过单个 GPU 的内存。一种常见做法是将嵌入表分片到不同设备上，这是一种模型并行的类型。为此，TorchRec 引入了其主要 API 称为DistributedModelParallel，或 DMP。与 PyTorch 的 DistributedDataParallel 类似，DMP 包装了一个模型以实现分布式训练。
安装
要求：python >= 3.7
在使用 TorchRec 时，我们强烈建议使用 CUDA（如果使用 CUDA：cuda >= 11.0）。
# install pytorch with cudatoolkit 11.3 conda install pytorch cudatoolkit=11.3 -c pytorch-nightly -y # install TorchRec pip3 install torchrec-nightly
概述
本教程将涵盖 TorchRec 的三个部分：nn.module EmbeddingBagCollection，DistributedModelParallel API 和数据结构KeyedJaggedTensor。
分布式设置
我们使用 torch.distributed 设置我们的环境。有关分布式的更多信息，请参见此tutorial。
在这里，我们使用一个 rank（colab 进程）对应于我们的 1 个 colab GPU。
import os import torch import torchrec import torch.distributed as dist os.environ["RANK"] = "0" os.environ["WORLD_SIZE"] = "1" os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29500" # Note - you will need a V100 or A100 to run tutorial as as! # If using an older GPU (such as colab free K80), # you will need to compile fbgemm with the appripriate CUDA architecture # or run with "gloo" on CPUs dist.init_process_group(backend="nccl")
从 EmbeddingBag 到 EmbeddingBagCollection
PyTorch 通过torch.nn.Embedding和torch.nn.EmbeddingBag来表示嵌入。EmbeddingBag 是 Embedding 的池化版本。
TorchRec 通过创建嵌入的集合来扩展这些模块。我们将使用EmbeddingBagCollection来表示一组 EmbeddingBags。
在这里，我们创建了一个包含两个嵌入包的 EmbeddingBagCollection（EBC）。每个表，product_table和user_table，由大小为 4096 的 64 维嵌入表示。请注意，我们最初将 EBC 分配到设备“meta”。这将告诉 EBC 暂时不要分配内存。
ebc = torchrec.EmbeddingBagCollection( device="meta", tables=[ torchrec.EmbeddingBagConfig( name="product_table", embedding_dim=64, num_embeddings=4096, feature_names=["product"], pooling=torchrec.PoolingType.SUM, ), torchrec.EmbeddingBagConfig( name="user_table", embedding_dim=64, num_embeddings=4096, feature_names=["user"], pooling=torchrec.PoolingType.SUM, ) ] )
DistributedModelParallel
现在，我们准备用DistributedModelParallel (DMP)包装我们的模型。实例化 DMP 将：
- 决定如何分片模型。DMP 将收集可用的“分片器”并提出一种最佳方式来分片嵌入表（即 EmbeddingBagCollection）的“计划”。
- 实际分片模型。这包括为每个嵌入表在适当设备上分配内存。
在这个示例中，由于我们有两个 EmbeddingTables 和一个 GPU，TorchRec 将两者都放在单个 GPU 上。
model = torchrec.distributed.DistributedModelParallel(ebc, device=torch.device("cuda")) print(model) print(model.plan)
使用输入和偏移查询普通的 nn.EmbeddingBag
我们使用input和offsets查询nn.Embedding和nn.EmbeddingBag。Input 是包含查找值的 1-D 张量。Offsets 是一个 1-D 张量，其中序列是每个示例要汇总的值的累积和。
让我们看一个例子，重新创建上面的产品 EmbeddingBag：
|------------| | product ID | |------------| | [101, 202] | | [] | | [303] | |------------|
product_eb = torch.nn.EmbeddingBag(4096, 64) product_eb(input=torch.tensor([101, 202, 303]), offsets=torch.tensor([0, 2, 2]))
使用 KeyedJaggedTensor 表示小批量
我们需要一个有效的表示，每个示例的每个特征中有任意数量的实体 ID 的多个示例。为了实现这种“不规则”表示，我们使用 TorchRec 数据结构KeyedJaggedTensor（KJT）。
让我们看看如何查找两个嵌入包“product”和“user”的集合。假设小批量由三个用户的三个示例组成。第一个示例有两个产品 ID，第二个没有，第三个有一个产品 ID。
|------------|------------| | product ID | user ID | |------------|------------| | [101, 202] | [404] | | [] | [505] | | [303] | [606] | |------------|------------|
查询应该是：
mb = torchrec.KeyedJaggedTensor( keys = ["product", "user"], values = torch.tensor([101, 202, 303, 404, 505, 606]).cuda(), lengths = torch.tensor([2, 0, 1, 1, 1, 1], dtype=torch.int64).cuda(), ) print(mb.to(torch.device("cpu")))
请注意，KJT 批量大小为batch_size = len(lengths)//len(keys)。在上面的例子中，batch_size 为 3。
将所有内容整合在一起，使用 KJT 小批量查询我们的分布式模型
最后，我们可以使用我们的产品和用户的小批量查询我们的模型。
结果查找将包含一个 KeyedTensor，其中每个键（或特征）包含一个大小为 3x64（batch_size x embedding_dim）的 2D 张量。
pooled_embeddings = model(mb) print(pooled_embeddings)
更多资源
有关更多信息，请参阅我们的dlrm示例，其中包括在 criteo terabyte 数据集上进行多节点训练，使用 Meta 的DLRM。
探索 TorchRec 分片
原文：
pytorch.org/tutorials/advanced/sharding.html译者：飞龙
本教程将主要涵盖通过EmbeddingPlanner和DistributedModelParallel API 对嵌入表进行分片方案的探索，并通过显式配置不同分片方案来探索嵌入表的不同分片方案的好处。
安装
要求：- python >= 3.7
我们强烈建议在使用 torchRec 时使用 CUDA。如果使用 CUDA：- cuda >= 11.0
# install conda to make installying pytorch with cudatoolkit 11.3 easier. !sudo rm Miniconda3-py37_4.9.2-Linux-x86_64.sh Miniconda3-py37_4.9.2-Linux-x86_64.sh.* !sudo wget https://repo.anaconda.com/miniconda/Miniconda3-py37_4.9.2-Linux-x86_64.sh !sudo chmod +x Miniconda3-py37_4.9.2-Linux-x86_64.sh !sudo bash ./Miniconda3-py37_4.9.2-Linux-x86_64.sh -b -f -p /usr/local
# install pytorch with cudatoolkit 11.3 !sudo conda install pytorch cudatoolkit=11.3 -c pytorch-nightly -y
安装 torchRec 还将安装FBGEMM，这是一组 CUDA 内核和 GPU 启用的操作，用于运行
# install torchrec !pip3 install torchrec-nightly
安装 multiprocess，它与 ipython 一起在 colab 中进行多进程编程
!pip3 install multiprocess
Colab 运行时需要以下步骤来检测添加的共享库。运行时在/usr/lib 中搜索共享库，因此我们复制在/usr/local/lib/中安装的库。这是非常必要的步骤，仅在 colab 运行时中。
!sudo cp /usr/local/lib/lib* /usr/lib/
**在此时重新启动您的运行时，以便看到新安装的软件包。**在重新启动后立即运行下面的步骤，以便 python 知道在哪里查找软件包。在重新启动运行时后始终运行此步骤。
import sys sys.path = ['', '/env/python', '/usr/local/lib/python37.zip', '/usr/local/lib/python3.7', '/usr/local/lib/python3.7/lib-dynload', '/usr/local/lib/python3.7/site-packages', './.local/lib/python3.7/site-packages']
分布式设置
由于笔记本环境，我们无法在此运行SPMD程序，但我们可以在笔记本内部进行多进程操作以模拟设置。用户在使用 Torchrec 时应负责设置自己的SPMD启动器。我们设置我们的环境，以便 torch 分布式基于通信后端可以工作。
import os import torch import torchrec os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29500"
PyTorch 2.2 中文官方教程（二十）（4）https://developer.aliyun.com/article/1482633
