TorchRec ShardedEmbeddingCollection
内部机制深度解析

1. 问题现象：为什么 modules() 看不到 lookups？

根本原因在于 ShardedEmbeddingModule 基类将三大核心子模块存储为普通 Python List，而不是 nn.ModuleList：

# torchrec/distributed/embedding_types.py:392-394
class ShardedEmbeddingModule(...):
    def __init__(self, ...):
        self._input_dists: List[nn.Module] = []   # ← 普通 list, 不是 nn.ModuleList!
        self._lookups: List[nn.Module] = []      # ← 普通 list, 不是 nn.ModuleList!
        self._output_dists: List[nn.Module] = [] # ← 普通 list, 不是 nn.ModuleList!

PyTorch 的 nn.Module.__setattr__ 只会自动注册以下类型到 _modules 字典：

• 直接赋值的 nn.Module 实例
• nn.ModuleList / nn.ModuleDict
• 通过 add_module() / register_module() 显式注册

普通 Python list 中的 nn.Module 不会被 PyTorch 追踪，因此 .modules()、.named_modules()、.parameters() 等方法都无法遍历到它们。

而 __repr__ 之所以能显示 lookups，是因为 ShardedEmbeddingModule 重写了 extra_repr()，手动遍历这些 list：

# torchrec/distributed/embedding_types.py:424-444
def extra_repr(self) -> str:
    rep = []
    rep.extend(loop("lookups", self._lookups))       # 手动遍历 list
    rep.extend(loop("_input_dists", self._input_dists))
    rep.extend(loop("_output_dists", self._output_dists))
    return "\n ".join(rep)

2. 完整模块层级结构

下图展示了 ShardedEmbeddingCollection 的完整内部结构。绿色为 PyTorch .modules() 可见部分，红色为不可见部分：

继承链

3. 为什么用 Plain List 而不是 nn.ModuleList？

4. TorchRec 重写 / Hook 了哪些 PyTorch 函数

以下是 TorchRec 在各个层级对 PyTorch nn.Module 标准方法的修改汇总：

4.1 ShardedEmbeddingCollection 层

4.2 GroupedEmbeddingsLookup 层

4.3 BatchedFusedEmbedding 层 (TBE Wrapper)

5. State Dict 完整工作流

5.1 保存路径: model.state_dict()

最终 state_dict 的 key 格式：embeddings.{table_name}.weight

5.2 加载路径: model.load_state_dict()

5.3 Hook 注册代码

# torchrec/distributed/embedding.py:1170-1175
self.register_state_dict_pre_hook(self._pre_state_dict_hook)
self._register_state_dict_hook(post_state_dict_hook)
self._register_load_state_dict_pre_hook(
    self._pre_load_state_dict_hook, with_module=True
)
self.register_load_state_dict_post_hook(_post_load_state_dict_hook)

6. 如何访问底层 FBGEMM TBE Module

方法一：直接遍历 _lookups（最常用）

sharded_ec: ShardedEmbeddingCollection = ...

for lookup in sharded_ec._lookups:
    # lookup 是 GroupedEmbeddingsLookup
    for emb_module in lookup._emb_modules:
        # emb_module 是 BatchedFusedEmbedding / BatchedDenseEmbedding / ...
        tbe = emb_module._emb_module
        # tbe 就是 SplitTableBatchedEmbeddingBagsCodegen (fbgemm TBE)
        print(type(tbe))
        # <class 'fbgemm_gpu...SplitTableBatchedEmbeddingBagsCodegen'>

方法二：获取所有 TBE 的通用辅助函数

def get_all_tbes(sharded_module):
    """从 ShardedEmbeddingCollection 或 ShardedEmbeddingBagCollection 提取所有 TBE"""
    tbes = []
    for lookup in sharded_module._lookups:
        # 处理 DDP 包裹
        while isinstance(lookup, DistributedDataParallel):
            lookup = lookup.module
        for emb_module in lookup._emb_modules:
            if hasattr(emb_module, '_emb_module'):
                tbes.append(emb_module._emb_module)
    return tbes

方法三：通过 named_parameters_by_table() 访问权重切片

# 按 table 获取权重 (不需要直接访问 TBE)
for lookup in sharded_ec._lookups:
    for table_name, tbe_slice in lookup.named_parameters_by_table():
        print(f"{table_name}: shape={tbe_slice.shape}, dtype={tbe_slice.dtype}")
        # tbe_slice 是 TableBatchedEmbeddingSlice(nn.Parameter)
        # 它是底层 TBE weight tensor 的一个 view

方法四：通过 DistributedModelParallel (DMP) 访问

# 如果你的模型被 DMP 包裹，sharded module 在 DMP 内部
dmp_model: DistributedModelParallel = ...

# 找到 sharded embedding collection
for name, module in dmp_model._dmp_wrapped_module.named_modules():
    if isinstance(module, ShardedEmbeddingCollection):
        for lookup in module._lookups:
            for emb in lookup._emb_modules:
                tbe = emb._emb_module  # fbgemm TBE

7. 关键类速查表

8. Embedding Weight 的双重身份

TorchRec 中 embedding weight 同时存在于两条路径，这是理解整个系统的关键：

_initialize_torch_state() 如何建立连接

# torchrec/distributed/embedding.py:919-929
for (table_name, tbe_slice) in lookup.named_parameters_by_table():
    # tbe_slice 是 TableBatchedEmbeddingSlice
    # 它底层是 TBE weight tensor 的一个 view
    self.embeddings[table_name].register_parameter("weight", tbe_slice)
    #          ↑ nn.ModuleDict 中的空 nn.Module
    #            现在有了 .weight 参数, 指向 TBE 内存

这样做的效果：

• state_dict() 遍历 self.embeddings 时能找到 .weight 参数
• parameters() 能正确返回所有 embedding weights
• 权重指向 TBE 的实际 GPU 内存，零拷贝
• 通过 post_state_dict_hook 在 checkpoint 时替换为 ShardedTensor/DTensor 格式