ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

DGL 多 GPU 分布式训练实战:基于 PyTorch DDP 与 GraphBolt 的 GraphSAGE 节点分类

DGL 多 GPU 分布式训练实战:基于 PyTorch DDP 与 GraphBolt 的 GraphSAGE 节点分类 人工智能机器学习深度学习图计算【免费下载链接】dglPython package built to ease deep learning on graph, on top of existing DL frameworks.项目地址https://gitcode.com/gh_mirrors/dg/dgl点击查看免费下载本篇技术指南以 DGL 官方多 GPU 示例为核心完整讲解如何在单机多卡环境下基于 PyTorchDistributedDataParallelDDP训练三层的 GraphSAGE-mean 节点分类模型。文章覆盖两种数据加载路径——经典dgl.dataloading.DataLoader与新一代graphbolt数据管道并深入源码剖析进程组初始化、DistributedItemSampler分布式采样、跨进程精度聚合dist.reduce加权平均、shared_tensor共享内存推理等关键机制帮助读者从能跑通进阶到理解为什么这么写。前置条件与环境要求硬件与软件要求多 GPU 训练的前提是机器上存在至少一张可用 GPU。示例脚本在入口处做了显式断言node_classification_sage.pyassert torch.cuda.is_available(), fMust have GPUs to enable multi-gpu training.此外还需要安装好 PyTorch含 CUDA 支持与 DGL 主库安装torchmetrics示例文档明确指定版本pip install torchmetrics0.11.4安装 OGBOpen Graph Benchmark数据包脚本通过ogb.nodeproppred.DglNodePropPredDataset拉取数据集若使用 GraphBolt 路径node_classification.py还需要确保 DGL 以启用 GraphBolt 的方式构建。目录结构与两个示例的关系examples/multigpu/目录下提供两条并行路线见 examples/multigpu文件数据加载方式对应单卡入门示例node_classification_sage.py经典dgl.dataloading.DataLoaderNeighborSamplerexamples/graphbolt/node_classification.py经典采样路线graphbolt/node_classification.py新一代dgl.graphbolt数据管道examples/graphbolt/node_classification.pyGraphBolt 路线两份脚本在训练流程上高度一致核心差异在于数据加载器与分布式采样器的实现下文将分别展开。经典路线node_classification_sage.py运行与参数详解启动命令与支持的数据集文档给出的最小启动命令为python3 node_classification_sage.py --dataset_name ogbn-products支持的数据集为ogbn-products与ogbn-arxiv。默认参数下文档记录的参考测试精度为* Test Accuracy of ogbn-products: ~0.7716 * Test Accuracy of ogbn-arxiv: ~0.6994说明以上精度为仓库文档在默认参数下的实测记录实际结果会因硬件、CUDA 版本、随机种子等因素略有浮动仅作参照。全部命令行参数一览脚本的argparse参数解析位于 node_classification_sage.py参数默认值取值/说明--modemixed训练模式可选mixedCPU-GPU 混合训练、puregpu纯 GPU 训练、benchmark基准测试模式--gpu0使用的 GPU 列表多卡用逗号分隔例如0,1,2,3--num_epochs10训练轮数--dataset_nameogbn-products数据集名可选ogbn-products、ogbn-arxiv--dataset_dirdataset数据集存放根目录--num_workers0DataLoader 工作进程数三种训练模式的含义--mode参数直接决定数据与图在 CPU/GPU 间的摆放位置其逻辑集中在 run 函数mixed默认g g.to(device if args.mode puregpu else cpu)——图驻留在 CPU 内存train_idx/val_idx移到 GPU并通过use_uva True启用UVAUnified Virtual Addressing统一虚拟寻址让 GPU 直接访问 CPU 内存中的图数据。这是大数据集下的推荐模式避免把整图塞进显存。puregpu整张图直接搬到 GPU 显存use_uva False。显存足够时速度最快。benchmark基准测试专用。此时索引不上 GPU、图也不搬移跳过if args.mode ! benchmark分支并使用NeighborSampler([10, 10, 10], fusedFalse)。注释说明这是为了规避 CUDA 运行错误对应 DGL issue #6697避免 fused 采样在 benchmark 场景下的已知问题。训练前的数据预处理主流程中除了加载数据集还包含几个容易被忽略但影响正确性的步骤node_classification_sage.pyAsNodePredDataset包装AsNodePredDataset(DglNodePropPredDataset(...))把 OGB 原始数据转换为 DGL 的节点预测标准格式统一提供train_idx/val_idx/test_idx与num_classes。g.create_formats_()显式建格式注释明确指出在多进程之前显式创建期望的图格式避免每个子进程重复创建并节省内存——这是多进程mp.spawn场景下的关键性能优化。ogbn-arxiv 特判g dgl.to_bidirected(g, copy_ndataTrue)再g dgl.add_self_loop(g)。由于 arxiv 是有向引文图需要转为无向并补自环才能用于 GraphSAGE 类模型。线程数限制os.environ[OMP_NUM_THREADS] str(mp.cpu_count() // 2 // nprocs)按进程数均分 CPU 线程避免多进程争抢资源。多进程 DDP 骨架run函数的四个关键步骤run是每个子进程的入口由mp.spawn(run, args(nprocs, devices, g, data, args), nprocsnprocs)启动每个 GPU 一个进程。源码中的注释node_classification_sage.py将其归纳为四步初始化进程组dist.init_process_group( backendnccl, # GPU 分布式训练使用 NCCL 后端 init_methodtcp://127.0.0.1:12345, world_sizenprocs, rankproc_id, )nccl是 PyTorch 官方推荐的 GPU 通信后端tcp://127.0.0.1:12345为单机多卡场景的简易 rendezvous 方式。若要多机扩展需替换为合适的主节点地址与端口。为子进程解包数据num_classes, train_idx, val_idx, test_idx data并将索引按模式搬到对应设备。实例化模型model SAGE(in_size, 256, num_classes).to(device)——输入维度取g.ndata[feat].shape[1]隐藏层 256输出为类别数。用 DDP 包装model DistributedDataParallel( model, device_ids[device], output_devicedevice )DDP 会在每个进程持有完整模型副本前向各自计算反向时通过 NCCL 自动同步梯度脚本注释opt.step() # Gradients are synchronized in DDP。这样每张卡上跑的是互斥的 mini-batch 子集但梯度全局一致等价于扩大了 batch size。GraphSAGE 模型定义模型使用 DGL 内置的dgl.nn.SAGEConv堆叠三层node_classification_sage.pyself.layers.append(dglnn.SAGEConv(in_size, hid_size, mean)) self.layers.append(dglnn.SAGEConv(hid_size, hid_size, mean)) self.layers.append(dglnn.SAGEConv(hid_size, out_size, mean)) self.dropout nn.Dropout(0.5)前向函数接收(blocks, x)即由邻居采样器产生的消息传递计算图列表每层一个 block逐层做SAGEConv → ReLU → Dropout最后一层不加激活。mean聚合器即 GraphSAGE-mean是 DGLSAGEConv支持的多种聚合方式之一。训练循环邻居采样、评估与跨进程精度聚合邻居采样器与 DataLoader训练与验证共用同一个NeighborSamplernode_classification_sage.pysampler NeighborSampler( [10, 10, 10], prefetch_node_feats[feat], prefetch_labels[label], )[10, 10, 10]三层网络每层采样 10 个邻居即 fanout 与层数严格对应prefetch_node_feats[feat]、prefetch_labels[label]启用 DGL 的数据预取在采样阶段把特征与标签一并取回减少训练循环内的设备间传输DataLoader 参数中use_ddpTrue是关键它让 DGL 在内部按world_size切分train_idx/val_idx使每个进程拿到互斥的子集源码注释use_ddpTrue, # To split the set for each processuse_uvause_uva在mixed模式下启用 UVA 直接读取 CPU 上的图训练 batch size 为 1024shuffleTrue验证集shuffleTrue验证顺序不影响精度仅影响聚合顺序。训练循环主体每轮 epoch 遍历 train dataloadernode_classification_sage.pyx blocks[0].srcdata[feat] # 第一层计算图的源节点特征 y blocks[-1].dstdata[label].to(torch.int64) # 最后一层目标节点标签 y_hat model(blocks, x) loss F.cross_entropy(y_hat, y) opt.zero_grad(); loss.backward(); opt.step()优化器为 Adamlr0.001, weight_decay5e-4。核心难点跨进程验证精度聚合每个进程只看到全局数据的一部分验证精度也必须先算局部、再全局归约。示例使用torch.distributed.reducenode_classification_sage.pyacc ( evaluate(device, model, g, num_classes, val_dataloader).to(device) / nprocs ) dist.reduce(tensoracc, dst0) if proc_id 0: print(fEpoch {epoch:05d} | Loss {total_loss / (it 1):.4f} | fAccuracy {acc.item():.4f} | Time {t1 - t0:.4f})原理每个进程先算出自己的验证精度除以进程数后用dist.reduce默认ReduceOp.SUM归约到 0 号进程0 号进程得到的就是全局平均精度。这种先除 nprocs 再 SUM的写法等价于加权平均前提是各进程样本量相等——若各进程样本数不等则需要像 GraphBolt 路线那样做按样本数加权的聚合见下文。源码注释也提示torch.distributed还提供其他归约算子如ReduceOp.MAX/MIN可按需选用。测试阶段layerwise_infer与共享内存推理测试阶段不再做邻居采样训练而是逐层全邻居推理node_classification_sage.pypred model.module.inference(g, device, batch_size, use_uva) pred pred[nid] labels g.ndata[label][nid].to(pred.device)SAGE.inference的实现node_classification_sage.py包含两个工程要点逐层处理对每一层使用MultiLayerFullNeighborSampler(1, prefetch_node_feats[h])即每层只看一跳全邻居len(blocks) 1配合大 batch默认2**10分块遍历全图避免整图前向的显存峰值。共享内存输出张量y shared_tensor((g.num_nodes(), self.hid_size if l ! len(self.layers) - 1 else self.out_size)) ... y[output_nodes] h.to(y.device, non_blockingTrue) dist.barrier()shared_tensor来自dgl.multiprocessing其实现位于 python/dgl/multiprocessing/pytorch.py它在 0 号进程创建一块共享内存张量并通过torch.distributed.broadcast把共享内存标识广播给其余进程使所有进程都能读写同一块宿主内存。这样每张卡算完自己那批节点的输出后直接non_blockingTrue写入共享张量配合页锁定内存加速拷贝全部写完后dist.barrier()同步保证 0 号进程读取时数据完整最后一层输出y取出测试节点子集pred[nid]计算MF.accuracy仅在proc_id 0打印最终测试精度。该设计既规避了多进程各自持有整图输出的显存开销也省去了全量特征的进程间通信。GraphBolt 路线graphbolt/node_classification.pyGraphBolt 是 DGL 新一代的数据加载与采样框架提供更高效、可组合的 DataPipe 风格 API。启动命令python node_classification.py --gpu0,1支持数据集在--dataset参数中限定为ogbn-arxiv、ogbn-products、ogbn-papers100M三种node_classification.py通过gb.BuiltinDataset(args.dataset).load()直接加载。BuiltinDataset的定义见 python/dgl/graphbolt/impl/ondisk_dataset.py它从远端存储下载内置数据集并封装为OnDiskDataset。参数表参数默认值说明--gpu0GPU 列表逗号分隔如0,1,2,3--epochs10训练轮数--lr0.001学习率--batch-size1024batch 大小--fanout10,10,10邻居采样 fanout长度须与模型层数一致--num-workers0DataLoader 工作进程数--gpu-cache-size0GPU 特征缓存容量字节大于 0 时启用gb.gpu_cached_feature--datasetogbn-products数据集三选一--modepinned-cuda数据放置/训练设备组合可选cpu-cuda、pinned-cuda、cuda-cuda三种存储模式与 GPU 特征缓存--mode由args.storage_device, _ args.mode.split(-)解析出数据存放位置node_classification.pycpu-cuda图与特征放在 CPU 内存模型在 GPUDataLoader 末尾调用datapipe.copy_to(device)把 mini-batch 搬到 GPU。pinned-cuda默认调用graph.pin_memory_()、feature.pin_memory_()把图与特征固定在页锁定pinned内存配合overlap_fetchTrue、asynchronousTrue让采样与特征预取可与 GPU 计算重叠是官方推荐的折中方案。cuda-cudadataset.graph.to(cuda)图与特征整体驻留显存速度最快但显存占用最高。当--gpu-cache-size 0且存储设备不是 CUDA 时还会用gb.gpu_cached_feature为节点特征建立 GPU 侧缓存实现大特征集 部分显存缓存的混合策略。create_dataloader分布式数据管道的组装这是 GraphBolt 路线的核心差异点node_classification.py整条管道由若干 DataPipe 串联datapipe gb.DistributedItemSampler( item_setitemset, batch_sizeargs.batch_size, drop_lastis_train, shuffleis_train, drop_uneven_inputsis_train, ) if args.storage_device ! cpu: datapipe datapipe.copy_to(device) datapipe datapipe.sample_neighbor( graph, args.fanout, overlap_fetchargs.storage_device pinned, asynchronousargs.storage_device ! cpu, ) datapipe datapipe.fetch_feature(features, node_feature_keys[feat]) if args.storage_device cpu: datapipe datapipe.copy_to(device) dataloader gb.DataLoader(datapipe, args.num_workers)其中gb.DistributedItemSampler是实现多卡数据切分的核心类见 python/dgl/graphbolt/item_sampler.py其关键行为与参数如下item_set当前数据集train_set/validation_set/test_setbatch_size每个 mini-batch 的样本数drop_last是否丢弃最后一个不满的 batch训练时为 Trueshuffle是否在采样前打乱训练时为 Truedrop_uneven_inputs是否强制各 rank 的 batch 数一致。其 docstring 明确说明当 DDP 训练中某 rank 输入更少时程序可能挂起或报错drop_uneven_inputsTrue通过丢弃多余 batch 解决同时也推荐使用 PyTorch 官方的 Join Context Manager 方案脚本在训练循环中即采用后者见下文seed可复现打乱的随机种子不指定则自动生成。关于该类的行为docstring 给出了 4 个组合示例15 个样本、4 个 replica、batch_size2默认drop_uneven_inputsFalse时Replica#3只有一个 batch[12,13]与[14]被分成两个其中最后一个不足drop_lastTrue时Replica#3只剩[12,13]而drop_uneven_inputsTrue会保证四个 replica 的 batch 数一致。它通过calculate_range将world_size、rank、num_workers、worker_id一起纳入索引区间划分torch.Generator().manual_seed(self._seed self._epoch)实现跨 epoch 可复现的洗牌。sample_neighbor、fetch_feature、copy_to则分别完成邻居采样、特征抓取与设备搬运overlap_fetch/asynchronous在 pinned 模式下启用异步流水线最大化数据加载与 GPU 计算的并行度。训练循环Join 上下文与加权聚合GraphBolt 版本在训练循环外包了一层Join([model])node_classification.pywith Join([model]): for data in (tqdm.tqdm(train_dataloader) if rank 0 else train_dataloader): x data.node_features[feat] y data.labels blocks data.blocks y_hat model(blocks, x) loss F.cross_entropy(y_hat, y) optimizer.zero_grad(); loss.backward(); optimizer.step() total_loss loss.detach() * y.size(0) num_train_items y.size(0)源码注释node_classification.py解释DDP 要求所有 rank 输入数一致否则可能报错或挂起PyTorch 的Join上下文管理器允许先跑完的 rank 等待其他 rankdrop_uneven_inputsTrue是备选方案。损失按样本数加权累计供后续聚合使用。每个 epoch 结束后通过自定义的weighted_reduce做加权平均node_classification.pydef weighted_reduce(tensor, weight, dst0): dist.reduce(tensortensor, dstdst) # 默认 ReduceOp.SUM weight torch.tensor(weight, devicetensor.device) dist.reduce(tensorweight, dstdst) return tensor / weight由于各 GPU 处理的样本数可能不同drop_last/drop_uneven_inputs的影响直接取平均会失真因此对损失×样本数与样本数分别做 SUM 归约后相除得到精确的全局均值total_loss weighted_reduce(total_loss, num_train_items) acc weighted_reduce(acc * num_val_items, num_val_items)torch.cuda.synchronize()在计时前强制同步保证 epoch 耗时测量准确最终rank 0打印Epoch | Average Loss | Accuracy | Time测试集同样用weighted_reduce聚合后输出Test Accuracy。进程管理与共享策略GraphBolt 版本在mp.spawn前额外设置了共享内存策略node_classification.pyos.environ[OMP_NUM_THREADS] str(mp.cpu_count() // 2 // world_size) mp.set_sharing_strategy(file_system) mp.spawn(run, args(world_size, args, devices, dataset), nprocsworld_size, joinTrue)mp.set_sharing_strategy(file_system)把进程间张量共享改为基于文件系统的策略避免fork场景下共享内存段的复制开销与冲突joinTrue让主进程等待所有子进程结束后再退出确保日志与清理顺序正确。两条路线的对比与选型建议维度经典路线node_classification_sage.pyGraphBolt 路线graphbolt/node_classification.py数据加载dgl.dataloading.DataLoaderNeighborSamplergb.DataLoaderDistributedItemSampler DataPipe 链分布式切分use_ddpTrue由 DGL 内部切分DistributedItemSampler显式按 rank/world_size 切分数据模式mixedUVA/puregpu/benchmarkcpu-cuda/pinned-cuda/cuda-cuda跨进程聚合acc / nprocs后dist.reduceSUMweighted_reduce按样本数加权归约不均匀输入处理依赖 DGL 内部切分保证均匀Join([model])或drop_uneven_inputsTrue特征缓存无--gpu-cache-size启用gb.gpu_cached_feature推理阶段SAGE.inferenceshared_tensor共享内存同左SAGE.inference逻辑选型建议数据集规模大、显存紧张优先mixed/pinned-cuda模式让图与特征留在 CPU/页锁定内存用 UVA 或异步预取喂给 GPU数据集可整体入显存且追求吞吐puregpu/cuda-cuda模式追求采样与预取流水线极致性能、或需要papers100M级别数据直接采用 GraphBolt 路线其 DataPipe 组合更灵活且自带 GPU 特征缓存选项无论哪条路线fanout[10,10,10]都必须与模型层数一致否则采样得到的 block 数量与SAGE.forward中的层循环不匹配。常见问题与排错要点CUDA 不可用脚本入口断言torch.cuda.is_available()无 GPU 环境会直接报错退出这是多 GPU 训练的前提。NCCL 初始化失败/端口冲突默认tcp://127.0.0.1:12345若端口被占用需修改init_method多机训练需替换为主节点的可达地址。验证/测试精度聚合错误经典路线先除 nprocs 再 SUM隐含各进程样本数相等若自定义了数据切分导致样本数不等应改用 GraphBolt 的weighted_reduce加权方案。训练挂起或不收敛优先检查各 rank 输入是否均匀drop_last/drop_uneven_inputs/Join以及OMP_NUM_THREADS线程数限制是否设置。显存不足mixed模式下shared_tensor把推理中间结果放在宿主内存训练侧可调小 batch size 或启用--gpu-cache-size特征缓存。延伸阅读经典单卡采样入门examples/graphbolt/node_classification.py采样版单卡对偶示例本目录另一份单卡对照examples/multigpu/README.md本文所依据的官方文档原文DGL 分布式数据加载实现python/dgl/dataloadingDataLoader与use_ddp机制shared_tensor共享内存实现python/dgl/multiprocessing/pytorch.pyGraphBolt 采样器源码python/dgl/graphbolt/item_sampler.py 与内置数据集封装 python/dgl/graphbolt/impl/ondisk_dataset.py多卡扩展参考examples/multigpu 同级的 examples/distributed 目录提供跨机分布式训练示例赞分享人工智能机器学习深度学习图计算【免费下载链接】dglPython package built to ease deep learning on graph, on top of existing DL frameworks.项目地址https://gitcode.com/gh_mirrors/dg/dgl点击查看免费下载相关推荐基于 PyTorch 的 DDP 分布式训练实战用 torchrun 编排多节点多卡应用基于 PyTorch 的 DDP 分布式训练实战用 torchrun 编排多节点多卡应用 本篇技术指南围绕当前仓库 distributed/ddp 示例系统示例工程人工智能深度学习基于DGL的GraphSAGE节点分类实现详解基于DGL的GraphSAGE节点分类实现详解 概述 本文将通过DGL Deep Graph Library 框架实现GraphSAGE算法在节点分类任务上的应人工智能机器学习深度学习图计算DGL分布式GraphSAGE训练实战指南DGL分布式GraphSAGE训练实战指南 概述 本文将详细介绍如何使用DGLDeep Graph Library框架进行分布式GraphSAGE模型训练。人工智能机器学习深度学习图计算创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表