DGL项目中使用ARGO提升多核CPU上的GNN训练性能-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_01179/article/details/148391656

DGL项目中使用ARGO提升多核CPU上的GNN训练性能

dgl Python package built to ease deep learning on graph, on top of existing DL frameworks. 项目地址: https://gitcode.com/gh_mirrors/dg/dgl

背景介绍

在深度学习领域，图神经网络(GNN)因其强大的图数据处理能力而广受关注。然而，当在多核CPU环境下训练GNN模型时，我们常常会遇到一个性能瓶颈：当核心数超过16个时，训练性能不再提升。DGL项目中的ARGO运行时系统正是为解决这一问题而设计。

ARGO通过智能的资源分配和核心绑定技术，能够有效利用多核CPU的计算能力。实验数据显示，在Xeon 8380H和Xeon 6430L处理器上，ARGO分别能带来4.30倍和3.32倍的训练加速。

ARGO系统安装

要使用ARGO系统，首先需要安装其依赖的scikit-optimize库，该库用于自动调参优化。安装方式有两种：

conda install -c conda-forge "scikit-optimize>=0.9.0"

或者

pip install scikit-optimize>=0.9

在DGL程序中启用ARGO

下面我们将详细介绍如何在现有的DGL GNN训练程序中启用ARGO优化。

第一步：导入必要模块

在程序开头添加以下导入语句：

import os
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
import torch.multiprocessing as mp
from argo import ARGO

确保argo.py文件与你的程序在同一目录下。

第二步：设置PyTorch分布式训练

修改训练函数，初始化分布式环境并包装模型：

def train(...):
    dist.init_process_group('gloo', rank=rank, world_size=world_size)
    model = SAGE(...)
    model = DistributedDataParallel(model)
    ...

在主程序中添加分布式设置：

os.environ['MASTER_ADDR'] = '127.0.0.1'
os.environ['MASTER_PORT'] = '29501'
mp.set_start_method('fork', force=True)
train(args, device, data)

第三步：初始化并运行ARGO

runtime = ARGO(n_search=15, epoch=args.num_epochs, batch_size=args.batch_size)
runtime.run(train, args=(args, device, data))

参数说明：

n_search: 搜索次数，建议设置在15-45之间以获得最佳性能
epoch: 训练总轮数
batch_size: 批处理大小

第四步：修改训练函数参数

扩展训练函数的参数列表：

def train(args, device, data, rank, world_size, comp_core, load_core, counter, b_size, ep):

新增参数用于ARGO的核心绑定和批处理调整。

第五步：修改数据加载器

调整数据加载器以支持ARGO优化：

dataloader = dgl.dataloading.DataLoader(
    g,
    train_nid,
    sampler,
    batch_size=b_size,
    shuffle=True,
    drop_last=False,
    num_workers=len(load_core),
    use_ddp=True)

第六步：启用CPU亲和性

在训练循环前添加核心绑定：

with dataloader.enable_cpu_affinity(loader_cores=load_core, compute_cores=comp_core):
    for epoch in range(ep):
        ...

第七步：模型保存与加载

添加模型检查点功能：

PATH = "model.pt"
if counter[0] != 0:
    checkpoint = th.load(PATH)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    epoch = checkpoint['epoch']
    loss = checkpoint['loss']

with dataloader.enable_cpu_affinity(...):
    ...

dist.barrier()
if rank == 0:
    th.save({
        'epoch': counter[0],
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss,
    }, PATH)