迁移单卡脚本为多卡脚本

本节为用户提供了将单卡训练脚本修改为多卡训练脚本的核心步骤指导,以自动混合精度(AMP)章节中的DDP场景(one NPU per process)的代码为样例进行介绍。

  1. 在主函数中添加如下代码。

    1. 在shell脚本中循环传入local_rank变量作为指定的device。
      local_rank = int(os.environ["LOCAL_RANK"])
    2. 用local_rank自动获取device号。
      device = torch.device('npu', local_rank)
    3. 初始化,将通信方式设置为hccl。
      torch.distributed.init_process_group(backend="hccl",rank=local_rank) 
    4. 在初始化时确定当前的device。
      torch_npu.npu.set_device(device)

  2. 在获取训练数据集后,设置train_sampler。

    train_sampler = torch.utils.data.distributed.DistributedSampler(train_data)

  3. 定义模型后,开启DDP模式。

    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)

  4. 将数据加载器train_dataloader与train_sampler相结合。

    train_dataloader = DataLoader(dataset = train_data, batch_size=batch_size, sampler = train_sampler)

  5. 如果用户要进行多机多卡训练,请前往准备多机多卡训练进行必要的配置,再参考拉起多卡分布式训练选择拉起分布式训练的方式、进行必要的修改和执行对应的拉起训练命令。如果用户要进行单机多卡训练,请直接前往拉起多卡分布式训练选择拉起分布式训练的方式、进行必要的修改和执行对应的拉起训练命令。