快速开始

介绍

本系统主要面向昇腾平台上的模型和算子开发者，提供便携易用的多机多卡内存访问方式，方便用户开发在卡间同步数据，加速通信或通算融合类算子开发。

软件架构

共享内存库接口主要分为host和device接口部分：

host侧接口提供初始化、内存管理、通信域管理以及同步功能。
device侧接口提供内存访问、同步以及通信域管理功能。

目录结构说明

├── 3rdparty // 依赖的第三方库
├── docs     // 文档
├── examples // 使用样例
├── include  // 头文件
├── scripts  // 相关脚本
├── src      // 源代码
└── tests    // 测试用例

软件硬件配套说明

硬件型号支持
- Atlas 800I A2/A3 系列产品
- Atlas 800T A2/A3 系列产品
平台：aarch64/x86
配套软件：驱动固件 Ascend HDK 25.0.RC1.1、 CANN 8.2.RC1及之后版本。（参考《CANN软件安装指南》安装CANN开发套件包以及配套固件和驱动）
cmake >= 3.19
GLIBC >= 2.28

快速上手

设置CANN环境变量

# root用户安装（默认路径）
source /usr/local/Ascend/ascend-toolkit/set_env.sh

共享内存库编译
编译共享内存库，设置共享内存库环境变量：
```
cd shmem
bash scripts/build.sh
source install/set_env.sh
```

run包使用
软件包名为：SHMEM_{version}_linux-{arch}.run
其中，{version}表示软件版本号，{arch}表示CPU架构。
安装run包（需要依赖cann环境）

chmod +x 软件包名.run # 增加对软件包的可执行权限
./软件包名.run --check # 校验软件包安装文件的一致性和完整性
./软件包名.run --install # 安装软件，可使用--help查询相关安装选项

出现提示xxx install success!则安装成功

shmem 默认开启tls通信加密。如果需要关闭，需要调用接口主动关闭：

int32_t ret = shmem_set_conf_store_tls(false, NULL, 0);

具体细节详见安全声明章节

执行一个样例matmul_allreduce算子。
1.在shmem/目录编译:

bash scripts/build.sh -examples

2.在shmem/examples/matmul_allreduce目录执行demo:

bash scripts/run.sh -ranks 2 -M 1024 -K 2048 -N 8192

注：example及其他样例代码仅供参考，在生产环境中请谨慎使用。

功能自测用例

bash scripts/build.sh -uttests
bash scripts/run.sh

run.sh脚本提供-ranks -ipport -test_filter等参数自定义执行用例的卡数、ip端口、gtest_filter等

例

# 8卡，ip:port 127.0.0.1:8666，运行所有*Init*用例
bash scripts/run.sh -ranks 8 -ipport "tcp://127.0.0.1:8666" -test_filter Init

python侧test用例

python接口API列表

在scripts目录下编译的时候，带上build python的选项

bash build.sh -python_extension

在install目录下，source环境变量

source set_env.sh

在src/python目录下，进行setup，获取到wheel安装包

python3 setup.py bdist_wheel

在src/python/dist目录下，安装wheel包

pip3 install shmem-xxx.whl --force-reinstall

设置是否开启TLS认证，默认开启，若关闭TLS认证，请使用如下接口

import shmem as shm
shm.set_conf_store_tls(False, "")   # 关闭tls认证

import shmem as shm
tls_info = "xxx"
shm.set_conf_store_tls(True, tls_info)      # 开启TLS认证

进入src/python目录，在该目录下使用torchrun运行测试demo

torchrun --nproc-per-node=k test.py # k为想运行的ranksize

看到日志中打印出“test.py running success!”即为demo运行成功

unique id 初始化方式¶

注：使用unique id的接口初始化，需要手动配置环境变量SHMEM_UID_SESSION_ID或者SHMEM_UID_SOCK_IFNAM，同时配置时只读SHMEM_UID_SESSION_ID SHMEM_UID_SESSION_ID配置示例： SHMEM_UID_SOCK_IFNAM=127.0.0.1:1234 SHMEM_UID_SOCK_IFNAM=[6666:6666:6666:6666:6666:6666:6666:6666]:886 SHMEM_UID_SOCK_IFNAM配置示例： SHMEM_UID_SOCK_IFNAM=enpxxxx:inet4 取ipv4 SHMEM_UID_SOCK_IFNAM=enpxxxx:inet6 取ipv6 不配置默认取eth:inet4

python初始化例子

import shmem as ash

# xxx

uid = ash.shmem_get_unique_id()
ret = ash.shmem_init_using_unique_id(rank, world_size, mem_size, uid)

# xxx

准备以上启动代码init.py后，使用torchrun --nproc-per-node 8 init.py，其中进程数 8 根据实际需要改动，更详细的例子，请参考unique_id_test.py文件。

c++初始化例子

shmem_uniqueid_t uid;
shmem_init_attr_t *attr;
int ret = shmem_get_uniqueid(&uid);
ret = shmem_set_attr(my_rank, n_ranks, mem_size, nullptr, &attr); // 第4个参数是ip_port，当前场景传入nullptr
ret = shmem_set_attr_uniqueid_args(my_rank, n_ranks, &uid, attr);