TorchRec中自定义稀疏特征分片方案的技术实践

背景介绍

在分布式模型训练场景中，稀疏特征的处理是一个关键挑战。PyTorch生态中的TorchRec库提供了强大的分布式嵌入功能，允许用户对稀疏特征进行高效的分片和并行处理。本文将深入探讨如何通过TorchRec实现自定义的稀疏特征分片方案。

核心问题

在推荐系统等场景中，某些稀疏特征之间存在频繁的交互关系。例如，用户ID、用户历史行为和用户画像特征经常在模型中被联合使用。如果这些特征被分散在不同GPU上，会导致大量的跨设备通信开销。

解决方案

TorchRec提供了construct_module_sharding_planAPI，允许开发者精细控制各个嵌入表的分片方式。通过该API，我们可以将具有强关联性的特征强制分配到同一GPU设备上。

关键实现步骤

1. 构建基础模型：首先创建包含EmbeddingBagCollection的标准模型结构

2. 定义分片方案：使用table_wise分片策略，明确指定每个嵌入表的目标设备

module_sharding_plan = construct_module_sharding_plan(
    model.sparse_arch.embedding_bag_collection,
    per_param_sharding={
        "table_0": table_wise(rank=0),
        "table_1": table_wise(rank=0),
        "table_2": table_wise(rank=0),
        "table_3": table_wise(rank=1),
        "table_4": table_wise(rank=1),
    },
)

3. 分布式模型包装：将分片方案应用到DistributedModelParallel中

model = DistributedModelParallel(
    module=model,
    device=device,
    plan=ShardingPlan({
        "model.sparse_arch.embedding_bag_collection": module_sharding_plan
    }),
)

注意事项

1. FBGEMM依赖：TorchRec底层依赖于FBGEMM库进行高效稀疏计算。若出现CUDA后端不支持的错误，需要确保正确安装FBGEMM的GPU版本。

2. 性能权衡：虽然将关联特征放在同一设备能减少通信，但也可能导致负载不均衡。需要根据实际特征访问模式和模型结构进行权衡。

3. 混合分片策略：可以结合使用table-wise和row-wise等不同分片策略，实现更灵活的特征分布。

最佳实践建议

1. 对于频繁交互的特征组，采用相同的rank分配
2. 对于大型稀疏特征，考虑使用row-wise分片以平衡内存
3. 监控各GPU的内存使用和通信开销，持续优化分片方案

通过这种精细化的分片控制，开发者可以在分布式训练中获得更好的性能和扩展性，特别是在推荐系统、广告排序等稀疏特征密集的场景中。