六、分布式嵌入

文章目录

六、分布式嵌入
前言
一、先要配置torch.distributed环境
[二、Distributed Embeddings](#二、Distributed Embeddings)
- [2.1 EmbeddingBagCollectionSharder](#2.1 EmbeddingBagCollectionSharder)
- [2.2 ShardedEmbeddingBagCollection](#2.2 ShardedEmbeddingBagCollection)
三、Planner
总结

前言

我们已经使用了TorchRec 的主模块：EmbeddedBagCollection 。我们在上一节研究了它是如何工作的，以及数据在TorchRec 中是如何表示的。然而，我们还没有探索TorchRec 的主要部分之一，即分布式嵌入

一、先要配置torch.distributed环境

EmbeddingBagCollectionSharder 依赖于 PyTorch 的分布式通信库（torch.distributed）来管理跨进程/GPU 的分片和通信。

首先初始化分布式环境

python 复制代码

import torch.distributed as dist

# 初始化进程组
dist.init_process_group(
    backend="nccl",          # GPU 推荐 NCCL 后端, CPU就是 gloo
    init_method="env://",    # 从环境变量读取节点信息
    rank=rank,               # 当前进程的全局唯一标识（从 0 开始）
    world_size=world_size,   # 总进程数（总 GPU 数）
)

pg = dist.GroupMember.WORLD

设置环境变量（多节点训练时必须）

bash 复制代码

import torch.distributed as dist

# 初始化进程组
# 在每个节点上设置以下环境变量
export MASTER_ADDR="主节点IP"   # 如 "192.168.1.1"
export MASTER_PORT="66666"     # 任意未占用端口
export WORLD_SIZE=4            # 总 GPU 数
export RANK=0                  # 当前节点的全局 rank

二、Distributed Embeddings

先回顾一下我们上一节的EmbeddingBagCollection module

代码演示：

python 复制代码

print(ebc)
"""
EmbeddingBagCollection(
  (embedding_bags): ModuleDict(
    (product_table): EmbeddingBag(4096, 64, mode='sum')
    (user_table): EmbeddingBag(4096, 64, mode='sum')
  )
)
"""

2.1 EmbeddingBagCollectionSharder

策略制定者，决定如何分片。
决定如何将 EmbeddingBagCollection 的嵌入表（Embedding Tables ）分布到多个 GPU/节点。
核心功能：根据配置（如 ShardingType ）生成分片计划（Sharding Plan）

代码演示：

python 复制代码

from torchrec.distributed.embedding_types import ShardingType
from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder

# 定义分片器：指定分片策略（如按表分片）
sharder = EmbeddingBagCollectionSharder(
    sharding_type=ShardingType.TABLE_WISE.value,  # 每个表分配到一个 GPU
    kernel_type=EmbeddingComputeKernel.FUSED.value,  # 使用 fused 优化
)

关键参数
- sharding_type ：分片策略，如：
  - TABLE_WISE：整个表放在一个 GPU。
  - ROW_WISE：按行分片到多个 GPU。
  - COLUMN_WISE：按列分片（适用于超大表）。
- kernel_type：计算内核类型（如 FUSED 优化显存）

2.2 ShardedEmbeddingBagCollection

策略执行者，实际管理分片后的嵌入表
根据 EmbeddingBagCollectionSharder 生成的分片计划，实际管理分布在多设备上的嵌入表。
核心功能：在分布式环境中执行前向传播、梯度聚合和参数更新

代码演示：

python 复制代码

from torchrec.distributed.embeddingbag import ShardedEmbeddingBagCollection

# 根据分片器生成分片后的模块
sharded_ebc = ShardedEmbeddingBagCollection(
    module=ebc,        # 原始 EmbeddingBagCollection
    sharder=sharder,   # 分片策略
    device=device,     # 目标设备（如 GPU:0）
)

三、Planner

它可以帮助我们确定最佳的分片配置。
Planner能够根据嵌入表的数量和GPU的数量来确定最佳配置。事实证明，这很难手动完成，工程师必须考虑大量因素来确保最佳的分片计划。
TorchRec 在提供的这个Planner ，可以帮助我们：
- 评估硬件的内存限制
- 将基于存储器获取的计算估计为嵌入查找
- 解决数据特定因素
- 考虑其他硬件细节，如带宽，以生成最佳分片计划

演示代码：

python 复制代码

from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology

# 初始化Planner
planner = EmbeddingShardingPlanner(
    topology=Topology(  # 硬件拓扑信息
        world_size=4,  # 总 GPU 数
        compute_device="cuda",
        local_world_size=2,  # 单机 GPU 数
        batch_size=1024,  
    ),
    constraints={  # 可选约束（如强制某些表使用特定策略）
        "user_id": ParameterConstraints(sharding_types=[ShardingType.TABLE_WISE]),
    },
)


# 生成分片计划
plan = planner.collective_plan(ebc, [sharder], pg)

# 分片后的模型
from torchrec.distributed.embeddingbag import ShardedEmbeddingBagCollection

sharded_ebc = ShardedEmbeddingBagCollection(
    module=ebc,
    sharder=sharder,
    device=torch.device("cuda:0"),
    plan=plan,  # 应用自动生成的分片计划
)

总结

TorchRec中的分布式嵌入以及训练设置。