仁怀企业网站建设,购买域名一般多少钱,职高动漫设计毕业后干什么,无锡新区网站制作#x1f4e6; 第一部分#xff1a;TorchRec 实战教程
TorchRec是PyTorch的领域库#xff0c;专为大规模推荐系统设计。其核心是解决超大规模嵌入表在多GPU/多节点上的高效训练问题。
1. 安装与环境配置 首先安装TorchRec及其依赖。推荐使用CUDA环境以获得最佳性能。
# 1.… 第一部分TorchRec 实战教程TorchRec是PyTorch的领域库专为大规模推荐系统设计。其核心是解决超大规模嵌入表在多GPU/多节点上的高效训练问题。1. 安装与环境配置首先安装TorchRec及其依赖。推荐使用CUDA环境以获得最佳性能。# 1. 安装对应CUDA版本的PyTorch (以CUDA 12.1为例)pipinstalltorch --index-url https://download.pytorch.org/whl/cu121# 2. 安装FBGEMM GPU版本和TorchRecpipinstallfbgemm-gpu --index-url https://download.pytorch.org/whl/cu121 pipinstalltorchrec --index-url https://download.pytorch.org/whl/cu121# 3. 如果是纯CPU环境性能较低# pip uninstall fbgemm-gpu -y# pip install fbgemm-gpu-cpu# pip install torchrec2. 核心概念分片与并行理解以下两个关键模块是使用TorchRec的基础分片器Sharder定义如何将巨大的嵌入表切割并分布到不同设备上。TorchRec支持多种分片策略如按行row_wise、按表table_wise等。分布式模型并行DistributedModelParallel, DMP这是TorchRec最核心的高级API。它类似于PyTorch的DistributedDataParallel但专为封装已分片的稀疏模型部分嵌入表和稠密模型部分如MLP而设计。3. 实战构建一个分布式推荐模型下面通过一个简化的代码示例展示如何使用TorchRec的关键组件。importtorchimporttorch.nnasnnfromtorchrec.distributedimportDistributedModelParallelfromtorchrec.distributed.plannerimportEmbeddingShardingPlannerfromtorchrec.modules.embedding_configsimportEmbeddingBagConfigfromtorchrec.modules.embedding_modulesimportEmbeddingBagCollectionfromtorchrec.distributed.model_parallelimport(get_default_sharders,)# 1. 定义模型以最简单的稠密-稀疏交互为例classSimpleRecModel(nn.Module):def__init__(self,embedding_bag_collection):super().__init__()self.ebcembedding_bag_collection# 假设稀疏特征维度总和为512self.densenn.Linear(512,1)defforward(self,sparse_features):embeddingsself.ebc(sparse_features)# 获得稀疏特征嵌入向量concatenatedtorch.cat([embforembinembeddings.values()],dim1)returnself.dense(concatenated)# 2. 初始化分布式环境必须在代码最开头importtorch.distributedasdist dist.init_process_group(backendnccl)# GPU用NCCLCPU用gloolocal_rankint(os.environ[LOCAL_RANK])devicetorch.device(fcuda:{local_rank})# 3. 配置嵌入表embedding_configs[EmbeddingBagConfig(nameuser,embedding_dim128,num_embeddings100_0000,# 一百万用户feature_names[user_feature],),EmbeddingBagConfig(nameitem,embedding_dim128,num_embeddings50_0000,# 五十万物品feature_names[item_feature],),]# 4. 在CPU上实例化模型重要DMP会处理设备移动ebcEmbeddingBagCollection(tablesembedding_configs,devicetorch.device(cpu))modelSimpleRecModel(ebc)# 5. 使用分布式模型并行DMP包装模型# get_default_sharders() 提供了适用于常见嵌入模块的分片器modelDistributedModelParallel(modulemodel,devicedevice,shardersget_default_sharders(),# plannerEmbeddingShardingPlanner() # 可选的自动规划器用于生成优化的分片计划)# 6. 定义优化器TorchRec的优化器支持稀疏更新高效处理嵌入梯度fromtorchrec.optimimportapply_optimizer_in_backwardfromtorch.optimimportSGD# 为嵌入参数设置稀疏优化器apply_optimizer_in_backward(SGD,model.module.ebc.parameters(),{lr:0.1})# 为稠密参数设置标准优化器dense_optimizerSGD(model.module.dense.parameters(),lr0.01)# 此后在训练循环中前向传播、反向传播和优化器步骤与非分布式模型基本一致。# DMP会自动处理跨设备的梯度同步和稀疏参数的更新。⚙️ 第二部分FSDP 快速指南FSDP是PyTorch原生的分布式训练策略核心思想是将模型的参数、梯度和优化器状态全部分片存储在需要时再通过通信收集从而极大节省单卡显存。1. 基本使用模式以下是使用FSDP包装一个Transformer模型的典型代码importtorchimporttorch.nnasnnfromtorch.distributed.fsdpimportFullyShardedDataParallelasFSDPfromtorch.distributed.fsdp.wrapimportdefault_auto_wrap_policy# 1. 初始化分布式环境 (同上略)# ...# 2. 定义一个大的模型 (例如Transformer)modelnn.Transformer(d_model2048,nhead16,num_encoder_layers12,num_decoder_layers12)# 3. 定义自动包装策略按子模块分片my_auto_wrap_policydefault_auto_wrap_policy(transformer_layer_cls{nn.TransformerEncoderLayer,nn.TransformerDecoderLayer})# 4. 用FSDP包装模型fsdp_modelFSDP(model,auto_wrap_policymy_auto_wrap_policy,device_idtorch.cuda.current_device(),)# 5. 定义优化器FSDP会自动处理优化器状态分片optimizertorch.optim.Adam(fsdp_model.parameters(),lr1e-4)# 6. 训练循环与非分布式模型一致# FSDP会在前向传播时透明地收集所需参数并在反向传播后同步梯度和更新分片。 第三部分TorchRec 与 FSDP 核心对比这两种技术都是为“大模型”设计但目标完全不同。下表清晰地展示了两者的区别对比维度TorchRecFSDP (Fully Sharded Data Parallel)核心目标专门用于大规模推荐系统解决稀疏嵌入表的并行训练。通用的大规模稠密模型如LLM、CV大模型训练解决参数显存瓶颈。主要并行范式混合并行嵌入表常采用模型并行/张量并行切分稠密部分使用数据并行。增强的数据并行在数据并行的基础上对参数、梯度、优化器状态进行分片。优化核心嵌入表的分片策略行、列、表以及稀疏梯度的高效聚合与更新。通信与计算的重叠以及分片策略全分片、混合分片以平衡显存和通信开销。关键优势1. 原生支持超大规模嵌入十亿/万亿级。2. 为推荐系统提供专用原语如EmbeddingBagCollection。3. 优化器支持稀疏更新计算高效。1.通用性强几乎适用于任何PyTorch模型。2.显存节省显著是训练千亿参数大模型的标配技术。3.与PyTorch生态无缝集成。典型应用场景电商推荐、广告点击率CTR预估、社交网络推荐等具有海量稀疏特征的场景。大语言模型LLM预训练与微调、大规模视觉模型训练、稠密科学计算模型。关系互补。一个复杂模型可同时使用两者其稀疏嵌入部分用TorchRec分片而稠密神经网络部分用FSDP分片。 第四部分如何选择与后续建议如何选择如果你的模型核心是处理用户ID、商品ID等海量离散特征嵌入表参数占模型绝大部分请直接选择TorchRec。如果你的模型是Transformer、ResNet等稠密结构参数巨大但并非稀疏特征应选择FSDP。对于混合模型大嵌入表大稠密网络可以研究组合使用两者。