5. PyTorch+NCCL源码编译

Pretend ^^2024-06-30 20:12

系列文章

目录

前言

从源码编译PyTorch和NCCL，可以实现对NCCL源码进行修改以适应特定需求，并应用于实际的分布式训练中，本文基于torch 2.2.1和nccl 2.19.3描述了一个大致过程，并验证了源码更改的有效性。

一、本地环境

  • Ubuntu 22.04.4 LTS (GNU/Linux 6.5.0-35-generic x86_64)
  • cuda 11.8+ cudnn 8
  • python 3.11
  • torch v2.2.1+ nccl v2.19.3
  • NVIDIA GeForce RTX 4090 *2

二、安装cudnn

下载cudnn包之后打开

bash 复制代码 
cd cudnn-linux-x86_64-8.9.7.29_cuda11-archive
# 复制到指定目录
sudo cp ./include/cudnn*.h /usr/local/cuda/include
sudo cp ./lib/libcudnn* /usr/local/cuda/lib64

chmod a+r /usr/local/cuda/include/cudnn*.h
chmod a+r /usr/local/cuda/lib64/libcudnn*

确认已经安装cudnn，除了cudnn_version.h，务必检查同目录下也有cudnn_ops_infer.h文件

bash 复制代码 
cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

可以看到对应cudnn版本为8.9.7

三、使用pytorch自带NCCL库进行编译安装

1. 源码编译

使用 python setup.py 命令进行源码编译，develop 命令通常在开发过程中使用，以在"开发模式"中安装包，其中对源代码的更改会立即生效而无需重新安装。develop更改为install 就是直接安装。

bash 复制代码 
# 新建conda虚拟环境，取名为nccl2
conda create -n nccl2 python=3.11
conda activate nccl2

#下载v2.2.1 源码
git clone --branch v2.2.1 --recursive https://github.com/pytorch/pytorch
cd pytorch      # v2.2.1 

# 安装依赖包
pip install -r requirements.txt

#以开发模式安装torch，不使用系统nccl，而是torch自带的，位于third party目录下 
MAX_JOBS=32 USE_CUDA=1 USE_NCCL=1 USE_SYSTEM_NCCL=0 python setup.py develop
  • 如下图所示即为开始编译：
  • 中途报错如下：貌似是numpy相关的错误
  • pip show numpy | grep Version查看numpy 版本,为2.0.0
  • 估计是numpy版本太新，导致一些变量名更改，=> 对numpy降级，实测1.26.3 可行, 之后make clean ，再重新编译
  • 源码编译的过程可能比较久，编译成功后提示如下，说明已经成功安装torch

2. 查看版本和all_reduce测试

编译完毕，测试能否用torch，cuda，nccl以及识别出GPU。这里新建了一个version.py

python 复制代码 
# version.py
import torch

print("torch version",torch.__version__)
print(torch.cuda.is_available(), torch.distributed.is_nccl_available())
print("nccl version:",torch.cuda.nccl.version())
print("cuda version:", torch.version.cuda)       
        
cudnn_version = torch.backends.cudnn.version()
print("cuDNN version:", cudnn_version)
print(torch.cuda.device_count(), torch.cuda.get_device_name(0))

结果如下,可以看到troch和nccl的版本，检测到双卡等。

执行以下代码，新建test.py, 使用 nccl 作为通信后端，在一个gpu上测试分布式训练中张量的 all_reduce 操作。

python 复制代码 
#test.py

import os
import torch
import torch.distributed as dist

os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'
dist.init_process_group("nccl", rank=0, world_size=1)
x = torch.ones(6)

if torch.cuda.is_available():
    y = x.cuda()
    dist.all_reduce(y)
    print(f"cuda allreduce result: {y}")

结果如下：

四、 修改NCCL源代码并重新编译后测试，体现出源码更改

修改 pytorch-2.2.1/third_party/nccl/nccl/src/collectives.cc 文件后，重新编译,

原代码如下，文件内包含了all_gather,all_reduce等各个集合通信操作，29行开始是All_Reduce的相关定义

cpp 复制代码 
/*************************************************************************
 * Copyright (c) 2015-2023, NVIDIA CORPORATION. All rights reserved.
 *
 * See LICENSE.txt for license information
 ************************************************************************/

#include "argcheck.h" // Need some checks here since we access comm
#include "collectives.h"
#include "enqueue.h"
#include "nccl.h"

NCCL_API(ncclResult_t, ncclAllGather, const void* sendbuff, void* recvbuff, size_t sendcount,
    ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream);
ncclResult_t ncclAllGather(const void* sendbuff, void* recvbuff, size_t sendcount,
    ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream) {
  // Just pass the size of one message and not the total bytes sent/received.
  constexpr nvtxPayloadSchemaEntry_t AllGatherSchema[] = {
    {0, NVTX_PAYLOAD_ENTRY_TYPE_SIZE, "Message size [bytes]"}
  };
  size_t msgsize = sendcount * ncclTypeSize(datatype);
  NVTX3_FUNC_WITH_PARAMS(AllGather, AllGatherSchema, msgsize)

  struct ncclInfo info = { ncclFuncAllGather, "AllGather",
    sendbuff, recvbuff, sendcount, datatype, ncclSum, 0, comm, stream, /* Args */
    ALLGATHER_CHUNKSTEPS, ALLGATHER_SLICESTEPS };
  return ncclEnqueueCheck(&info);
}

NCCL_API(ncclResult_t, ncclAllReduce, const void* sendbuff, void* recvbuff, size_t count,ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, cudaStream_t stream);
ncclResult_t ncclAllReduce(const void* sendbuff, void* recvbuff, size_t count,ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, cudaStream_t stream) 
{
  struct NvtxParamsAllReduce {
    size_t bytes;
    ncclRedOp_t op;
  };
  // Just pass the size of one message and not the total bytes sent/received.
  static constexpr nvtxPayloadSchemaEntry_t AllReduceSchema[] = {
    {0, NVTX_PAYLOAD_ENTRY_TYPE_SIZE, "Message size [bytes]"},
    {0, NVTX_PAYLOAD_ENTRY_NCCL_REDOP, "Reduction operation", nullptr, 0,
      offsetof(NvtxParamsAllReduce, op)}
  };
  NvtxParamsAllReduce payload{count * ncclTypeSize(datatype), op};
  NVTX3_FUNC_WITH_PARAMS(AllReduce, AllReduceSchema, payload)

  struct ncclInfo info = { ncclFuncAllReduce, "AllReduce",
    sendbuff, recvbuff, count, datatype, op, 0, comm, stream, /* Args */
    ALLREDUCE_CHUNKSTEPS, ALLREDUCE_SLICESTEPS };
  return ncclEnqueueCheck(&info);
}

修改ncclAllReduce函数， 将内部全部注释掉，加一句 return ncclSystemError;

cpp 复制代码 
NCCL_API(ncclResult_t, ncclAllReduce, const void* sendbuff, void* recvbuff, size_t count,ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, cudaStream_t stream);

ncclResult_t ncclAllReduce(const void* sendbuff, void* recvbuff, size_t count,ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, cudaStream_t stream) 
{
  // struct NvtxParamsAllReduce {
  //   size_t bytes;
  //   ncclRedOp_t op;
  // };
  // // Just pass the size of one message and not the total bytes sent/received.
  // static constexpr nvtxPayloadSchemaEntry_t AllReduceSchema[] = {
  //   {0, NVTX_PAYLOAD_ENTRY_TYPE_SIZE, "Message size [bytes]"},
  //   {0, NVTX_PAYLOAD_ENTRY_NCCL_REDOP, "Reduction operation", nullptr, 0,
  //     offsetof(NvtxParamsAllReduce, op)}
  // };
  // NvtxParamsAllReduce payload{count * ncclTypeSize(datatype), op};
  // NVTX3_FUNC_WITH_PARAMS(AllReduce, AllReduceSchema, payload)

  // struct ncclInfo info = { ncclFuncAllReduce, "AllReduce",
  //   sendbuff, recvbuff, count, datatype, op, 0, comm, stream, /* Args */
  //   ALLREDUCE_CHUNKSTEPS, ALLREDUCE_SLICESTEPS };
  // return ncclEnqueueCheck(&info);
  return ncclSystemError;
}

每次修改pytorch中nccl源码,要使之生效需要进行重新编译，先删除原有编译文件再重新编译

bash 复制代码 
#删除原有nccl相关的
rm -r ./build/nccl*

#重新编译
MAX_JOBS=32 USE_CUDA=1 USE_NCCL=1 USE_SYSTEM_NCCL=0  python setup.py develop

#运行测试文件，看看有没有报错
python test.py

如图：报错ncclSystemError，体现出了源码的更改。

以后就可以按照这种方法修改nccl源码，使之与pytorch集成，将修改后的nccl应用于实际的分布式训练中了。

关于nccl源码及大致的总体流程，推荐一个大佬的文章，写的比较详细，令我受益匪浅。
https://blog.csdn.net/kidgin7439/category_11998768.html

相关推荐
我非夏日8 分钟前
基于Hadoop平台的电信客服数据的处理与分析③项目开发：搭建基于Hadoop的全分布式集群---任务9：HBase的安装和部署
大数据·hadoop·分布式·hbase·大数据技术开发
这不比博人传燃？1 小时前
actual combat 32 —— RabbitMQ
分布式·rabbitmq
小生凡一1 小时前
图解 Kafka 架构
分布式·架构·kafka
LightOfNight1 小时前
【后端面试题】【中间件】【NoSQL】ElasticSearch的优化方案2（减少字段、冷热分离、参数优化）
分布式·后端·elasticsearch·中间件·架构·nosql
InterestingFigure1 小时前
头条系统-05-延迟队列精准发布文章-概述&添加任务(db和redis实现延迟任务)、取消&拉取任务&定时刷新(redis管道、分布式锁setNx)...
数据库·redis·分布式·缓存
天才的白鸟2 小时前
mapreduce实现bean的序列化与反序列化
大数据·hadoop·分布式
Comedy_宁3 小时前
分布式NAS集群+ceph+CTDB
linux·分布式·ceph·nas·ctdb
LightOfNight3 小时前
【后端面试题】【中间件】【NoSQL】MongoDB的配置服务器、复制机制、写入语义和面试准备
分布式·后端·mongodb·中间件·面试·架构·nosql
JAMES费3 小时前
python机器人编程——用pytorch实现六轴机械臂的正向和逆向数值解算，及python算法解析
pytorch·python·机器人
xintaiideas3 小时前
深入理解 RabbitMQ、RocketMQ等常⽤的消息中间件进⾏消息的异步数据处理
分布式·rabbitmq·rocketmq
热门推荐
01Coze扣子平台完整体验和实践（附国内和国际版对比）02组基轨迹建模 GBTM的介绍与实现（Stata 或 R）03【经验分享】Ubuntu22.04安装微信(linux官方版)042025张宇考研数学，百度网盘视频课+36讲PDF讲义+真题05【立创EDA-PCB设计基础】6.布线铺铜实战及细节详解06yolov8实战第五天——yolov8+ffmpg实时视频流检测并进行实时推流——（推流，保姆教学）0751-2 万字长文，深度解读端到端自动驾驶的挑战和前沿08通达信神奇九转指标原理及选股公式，无未来函数，数字不消失09浏览器插件——ModHeader 的分享和学习10＜工控＞＜PLC＞汇川Easy521系列PLC与汇川SV630N伺服进行EtherCat通讯的相关配置及指令编写