摩尔线程S4000国产信创计算卡性能实战——Pytorch转译，多卡P2P通信与MUSA编程

06-02 1321阅读

简介

MTT S4000 是基于摩尔线程曲院 GPU 架构打造的全功能元计算卡，为千亿规模大语言模型的训练、微调和推理进行了定制优化，结合先进的图形渲染能力、视频编解码能力和超高清 8K HDR 显示能力，助力人工智能、图形渲染、多媒体、科学计算与物理仿真等复合应用场景的计算加速。

MTT S4000 全面支持大语言模型的预训练、微调和推理服务，MUSA 软件栈专门针对大规模集群的分布式计算性能进行了优化，适配主流分布式计算加速框架， 包括 DeepSpeed， Colossal AI，Megatron 等，支持千亿参数大语言模型的稳定预训练。

官方参数如下

摩尔线程S4000国产信创计算卡性能实战——Pytorch转译，多卡P2P通信与MUSA编程

运行环境

本次运行环境为AutoDL云中的镜像环境，系统环境如下

摩尔线程S4000国产信创计算卡性能实战——Pytorch转译，多卡P2P通信与MUSA编程

常用命令

显卡运行状态

输入如下命令

mthreads-gmi

即可查看当前显卡运行状态

摩尔线程S4000国产信创计算卡性能实战——Pytorch转译，多卡P2P通信与MUSA编程

查看当前GPU详细信息

输入

musaInfo

即可

摩尔线程S4000国产信创计算卡性能实战——Pytorch转译，多卡P2P通信与MUSA编程

查看当前运行环境版本

输入

musa_version_query

即可查看当前运行环境版本

摩尔线程S4000国产信创计算卡性能实战——Pytorch转译，多卡P2P通信与MUSA编程

Pytorch部分

转义

根据官网介绍，对于pytorch代码，只需要正确import torch_musa的拓展插件，并且将代码中的所有cuda->musa，将所有的nccl->mccl即可。

摩尔线程S4000国产信创计算卡性能实战——Pytorch转译，多卡P2P通信与MUSA编程

实测

作者使用豆包随机生成了一个测试allreduce的pytorch代码，代码如下，在经过上述转译后能正常运行

import os
import time
import argparse
import torch
import torch_musa
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    # 初始化MUSA分布式环境
    dist.init_process_group("mccl", rank=rank, world_size=world_size)
    torch.musa.set_device(rank)
def cleanup():
    dist.destroy_process_group()
def run_benchmark(rank, world_size, sizes, num_iters=100, warmup=20):
    setup(rank, world_size)
    for size in sizes:
        # 创建随机张量（使用MUSA设备）
        tensor = torch.rand(size, device=f'musa:{rank}')
        # 预热
        for _ in range(warmup):
            dist.all_reduce(tensor)
            torch.musa.synchronize()
        # 测量时间
        start_time = time.time()
        for _ in range(num_iters):
            dist.all_reduce(tensor)
            torch.musa.synchronize()
        end_time = time.time()
        # 计算统计信息
        total_time = end_time - start_time
        avg_time = total_time / num_iters
        size_mb = size * 4 / (1024 * 1024)  # float32是4字节
        bandwidth = (size_mb * world_size) / avg_time  # MB/s
        if rank == 0:
            print(f"张量大小: {size:,} 元素 ({size_mb:.2f} MB)")
            print(f"平均耗时: {avg_time * 1000:.2f} ms")
            print(f"带宽: {bandwidth / 1024:.2f} GB/s")
            print("-" * 50)
    cleanup()
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--sizes', type=int, nargs='+',
                        default=[1000, 10000, 100000, 1000000, 10000000, 100000000],
                        metavar='N',
                        help='测试的张量大小列表')
    parser.add_argument('--num-iters', type=int, default=100,
                        help='每个大小的迭代次数')
    parser.add_argument('--warmup', type=int, default=20,
                        help='预热迭代次数')
    args = parser.parse_args()
    world_size = torch.musa.device_count()
    if world_size != 4:
        raise ValueError("此脚本需要4个MUSA GPU，但发现 {} 个GPU".format(world_size))
    import torch.multiprocessing as mp
    mp.spawn(run_benchmark,
             args=(world_size, args.sizes, args.num_iters, args.warmup),
             nprocs=world_size,
             join=True)
if __name__ == "__main__":
    main()

摩尔线程S4000国产信创计算卡性能实战——Pytorch转译，多卡P2P通信与MUSA编程摩尔线程S4000国产信创计算卡性能实战——Pytorch转译，多卡P2P通信与MUSA编程

MUSA编程

p2p通信部分

代码参考

笔者按照英伟达cudasamples仓库中的p2pbandwidthtest 代码，cuda-samples/Samples/5_Domain_Specific/p2pBandwidthLatencyTest at master · NVIDIA/cuda-samples · GitHub

并且参考相应的musa event api与mublasapi

https://docs.mthreads.com/musa-sdk/musa-sdk-doc-online/api/mcc_um.zh-CN

编写了一个适用于摩尔线程的p2p通信检测验证程序

代码部分

#include 
#include 
#include   // 假设 MUSA 头文件
using namespace std;
const char *sSampleName = "P2P (Peer-to-Peer) GPU Bandwidth Latency Test";
typedef enum {
    P2P_WRITE = 0,
    P2P_READ  = 1,
} P2PDataTransfer;
typedef enum {
    CE = 0,
    SM = 1,
} P2PEngine;
P2PEngine p2p_mechanism = CE; // 默认使用 Copy Engine
// 错误检查宏
#define musaCheckError()               \
    {                                    \
        musaError_t e = musaGetLastError();                                             \
        if (e != musaSuccess) {          \
            printf("MUSA failure %s:%d: '%s'\n", __FILE__, __LINE__, musaGetErrorString(e)); \
            exit(EXIT_FAILURE);          \
        }                                \
    }
// 延迟内核
__global__ void delay(volatile int *flag, unsigned long long timeout_clocks = 10000000)
{
    // 等待应用程序通知我们它已经完成了实验的排队，或者超时并退出，允许应用程序继续执行
    long long int start_clock, sample_clock;
    start_clock = clock64();
    while (!*flag) {
        sample_clock = clock64();
        if (sample_clock - start_clock > timeout_clocks) {
            break;
        }
    }
}
// P2P 复制内核
__global__ void copyp2p(int4 *__restrict__ dest, const int4 *__restrict__ src, size_t num_elems) {
    size_t globalId = blockIdx.x * blockDim.x + threadIdx.x;
    size_t gridSize = blockDim.x * gridDim.x;
#pragma unroll 5 // 移除括号
    for (size_t i = globalId; i  

编译
 

参考mcc编译手册，此时代码中引用的库为musa_runtime,则编译是-l参数后跟musart
 

mcc p2p.mu -o p2p -lmusart

 

结果
 

可以看到p2p已经正确开启，但是延迟测试有问题，后续改进
 

摩尔线程S4000国产信创计算卡性能实战——Pytorch转译，多卡P2P通信与MUSA编程
 

基于musa编程的allreduce测试
 

代码参考
 

主要参考了NCCLtest中的allreduce部分逻辑
 

GitHub - NVIDIA/nccl-tests: NCCL Tests
 

并且参考了mublas api设计
https://docs.mthreads.com/musa-sdk/musa-sdk-doc-online/api/mublas_api
 

代码部分
 

#include 
#include 
#include 
#include "musa_runtime.h"
#include "mccl.h"
#include  // 必须包含此头文件
// 宏定义（所有标识符在此处声明）
#define MIN_SIZE_B       16ULL           // 最小测试尺寸（16字节）
#define MAX_SIZE_B  (4096ULL * 1024ULL * 1024ULL)  // 最大测试尺寸（4096MB）
#define STEP_FACTOR      2ULL           // 尺寸增长因子（每次翻倍）
#define WARMUP_ITERS       5             // 热身迭代次数
#define TEST_ITERS        20             // 测试迭代次数
#define ROOT_RANK        -1             // 根节点（-1表示全归约）
#define DATA_TYPE         mcclFloat      // 数据类型
#define REDUCTION_OP      mcclSum        // 归约操作
#define FLOAT_SIZE        sizeof(float)  // float类型字节数（4字节）
// 错误检查宏
#define MUSACHECK(cmd) do { \
    musaError_t err = cmd; \
    if (err != musaSuccess) { \
        printf("MUSA Error at %s:%d: %s\n", __FILE__, __LINE__, musaGetErrorString(err)); \
        exit(EXIT_FAILURE); \
    } \
} while(0)
#define MCCLCHECK(cmd) do { \
    mcclResult_t res = cmd; \
    if (res != mcclSuccess) { \
        printf("MCCL Error at %s:%d: %s\n", __FILE__, __LINE__, mcclGetErrorString(res)); \
        exit(EXIT_FAILURE); \
    } \
} while(0)
// 带宽计算函数
void calculate_bandwidth(size_t count, int type_size, double time_sec, double* alg_bw, double* bus_bw, int nranks) {
    if (time_sec  1) ? (2.0 * (nranks - 1)) / nranks : 1.0;
    *bus_bw = *alg_bw * factor;
}
int main(int argc, char* argv[]) {
    int nDev = 4;                  // 设备数量
    int devs[4] = {0, 1, 2, 3};     // 设备ID列表
    mcclComm_t comms[4];           // MCCL通信器
    musaStream_t streams[4];       // 流数组
    float** sendbuff = NULL;       // 发送缓冲区
    float** recvbuff = NULL;       // 接收缓冲区
    size_t current_size_b = MIN_SIZE_B;  // 当前测试尺寸（字节）
    double alg_bw, bus_bw;          // 算法带宽和总线带宽
    int test_wrong = 0;             // 错误计数
    // 初始化MCCL通信器
    MCCLCHECK(mcclCommInitAll(comms, nDev, devs));
    // 分配设备内存并创建流
    sendbuff = (float**)malloc(nDev * sizeof(float*));
    recvbuff = (float**)malloc(nDev * sizeof(float*));
    for (int i = 0; i  

编译
 

因为代码用了musa_runtime与mccl两个库，因此编译选项也会有所改变
 

mcc allreduce.mu -o allreduce -lmusart -lmccl

 

结果
 

不知道为什么结果测出来和用pytorch测出来结果相差不小，目测是因为musa event打点计时函数没使用正确（在p2p测试的自交中也有体现，不管什么情况都是50us左右），这个需要后续再看下
 

摩尔线程S4000国产信创计算卡性能实战——Pytorch转译，多卡P2P通信与MUSA编程
 

摩尔线程S4000国产信创计算卡性能实战——Pytorch转译，多卡P2P通信与MUSA编程
免责声明：我们致力于保护作者版权，注重分享，被刊用文章因无法核实真实出处，未能及时与作者取得联系，或有版权异议的，请联系管理员，我们会立即处理! 部分文章是来自自研大数据AI进行生成,内容摘自(百度百科,百度知道,头条百科,中国民法典,刑法,牛津词典,新华词典,汉语词典,国家院校,科普平台)等数据,内容仅供学习参考,不准确地方联系删除处理! 图片声明：本站部分配图来自人工智能系统AI生成,觅知网授权图片,PxHere摄影无版权图库和百度，360，搜狗等多加搜索引擎自动关键词搜索配图，如有侵权的图片，请第一时间联系我们。

相关阅读

目录[+]

取消
微信二维码
微信二维码
支付宝二维码