单机多卡场景
环境准备
- 请参考《Ascend Extension for PyTorch 配置与安装》手册,准备昇腾环境,包括硬件与OS,NPU驱动固件、CANN软件、PyTorch的安装。
- 根据实际需求准备模型环境,如conda、docker以及三方库依赖。分布式训练场景下,HCCL会使用Host服务器的部分端口进行集群信息收集,需要操作系统预留该部分端口。HCCL建议使用60000-60015端口,若通过环境变量HCCL_IF_BASE_PORT指定了Host网卡起始端口,则需要预留以该端口起始的16个端口。
- 若操作系统端口号预留仅需临时生效,可执行如下命令,临时生效的配置重启后会失效:
sysctl -w net.ipv4.ip_local_reserved_ports=60000-60015
- 若操作系统端口号预留需永久生效,可执行如下命令:
- 以root用户登录服务器,编辑“/etc/sysctl.conf”文件。
vim /etc/sysctl.conf
- 在“/etc/sysctl.conf”文件末尾加上net.ipv4.ip_local_reserved_ports=60000-60015,保存并退出。
- 执行如下命令使配置生效。
sysctl -p
- 以root用户登录服务器,编辑“/etc/sysctl.conf”文件。
- 若操作系统端口号预留仅需临时生效,可执行如下命令,临时生效的配置重启后会失效:
模型脚本配置
本节以适配样例(DDP场景)章节的代码为样例,介绍将单机单卡训练脚本修改为单机多卡训练脚本的核心步骤。
- 在主函数中添加如下代码。
- 在shell脚本中循环传入local_rank变量作为指定的device。
local_rank = int(os.environ["LOCAL_RANK"])
- 用local_rank自动获取device号。
device = torch.device('npu', local_rank)
- 初始化,将通信方式设置为hccl。
torch.distributed.init_process_group(backend="hccl",rank=local_rank)
- 在shell脚本中循环传入local_rank变量作为指定的device。
- 在获取训练数据集后,设置train_sampler。
train_sampler = torch.utils.data.distributed.DistributedSampler(train_data)
- 定义模型后,开启DDP模式。
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
- 将数据加载器train_dataloader与train_sampler相结合。
train_dataloader = DataLoader(dataset = train_data, batch_size=batch_size, sampler = train_sampler)
启动脚本配置
有5种脚本启动方式可拉起多卡训练:
- shell脚本方式(推荐)
- mp.spawn方式
- Python方式
- torchrun方式
- torch_npu_run方式(集群场景推荐):此方式是torchrun在大集群场景的改进版,提升集群建链性能。
附录拉起单机八卡训练中,以一个简单模型脚本为样例,展示了每种拉起方式脚本代码的修改方法以及各种拉起方式的适配方法,用户可以参考学习。

- 集合通信存在如下约束:
- 数据并行模式中不同device上执行的图相同。
- 针对Atlas 训练系列产品:allreduce和reduce_scatter仅支持int8、int32、float16和float32数据类型。
- 针对Atlas A2 训练系列产品:allreduce和reduce_scatter仅支持int8、int32、float16、float32和bfp16数据类型。
- 针对Atlas 训练系列产品,如果使用1台训练服务器(Server),要求实际参与集合通信的NPU数目只能为1/2/4/8,且0-3卡和4-7卡各为一个组网。使用2张卡或4张卡训练时,不支持跨组网创建设备集群。
- 针对Atlas 训练系列产品,Server集群场景下(即由集群管理主节点和一组训练服务器组成训练服务器集群),要求参与集合通信的NPU数目只能为1*n、2*n、4*n、8*n(其中n为参与训练的Server个数,上限为512)。且n为2的指数倍情况下,集群性能最好,建议用户优先采用此种方式进行集群组网。
- 针对Atlas A2 训练系列产品,Server集群场景(即由集群管理主节点和一组训练服务器组成训练服务器集群)下要求参与集合通信的NPU数目为(1~8)*n(其中n为参与训练的Server个数,上限为1152),其中,n为2的指数倍情况下,集群性能最好,建议用户优先采用此种方式进行集群组网。同时建议每个Sever中参与集合通信的NPU数量保持一致,若不一致,会造成性能裂化。
- 针对Atlas A2 训练系列产品,若用户准备进行2卡训练,可将8卡训练脚本进行改写,改为2卡训练脚本。可参见以下修改方法:
- 若8卡脚本的batchsize是单卡脚本的batchsize的8倍,则将8卡训练时的batch size和learning rate同时除以4,作为2卡训练时的batch size和learning rate。
- 如果使用for循环启动训练入口脚本,则将for循环的次数改为2次。
- world size或者rank size修改为2,并确保训练脚本中dist.init_process_group()中world_size参数为2。
- 如果有指定device list参数,且取值范围为0-7,则将其改为0-1。
- 一个Device对应执行一个训练进程,当前不支持多进程在同一个Device上进行训练。
父主题: 模型脚本与启动脚本配置