Skip to main content

分布式改造

如果在线开发这种单机模式训练速度较慢,则需要把训练代码进行分布式化改造以利用平台分布式多机训练的功能。

改造指南

Pytorch

可参考官方教程以及一些网络文章,大部分情况下只需要添加分布式API即可。

示例: 将deeplabv3的单机训练代码进行分布式化后的主要变化为:

  • 设置分布式训练任务之间的通信后端,推荐NCCL

    if dist.is_available():
    parser.add_argument('--backend', type=str, help='Distributed backend',
    choices=[dist.Backend.GLOO, dist.Backend.NCCL, dist.Backend.MPI],
    default=dist.Backend.NCCL)

    dist.init_process_group(backend=args.backend)
  • 获取注入的环境变量如WORLD_SIZE代表总共的训练节点数,RANK代表当前节点是第几个训练节点,这些环境变量会通过KubeFlow的training operator或是自定义的launcher自动注入,类似分布式任务的环境变量还有:

    #将NCCL_IB_DISABLE设置为1来禁止使用InfiniBand,转而使用 IP;如果网络接口不能被自动发现,则手工设置NCCL_SOCKET_IFNAME
    NCCL_IB_DISABLE=1
    #训练pod的名称
    HOSTNAME=pytorchjob-seg-1275-master-0
    NVIDIA_VISIBLE_DEVICES=GPU-46c4b4d9-74ef-4d5d-7067-583c2ba5f993
    KUBERNETES_PORT_443_TCP_PROTO=tcp
    KUBERNETES_PORT_443_TCP_ADDR=10.43.0.1
    KUBERNETES_PORT=tcp://10.43.0.1:443
    #launcher传入的工作目录
    PWD=/app
    HOME=/root
    KUBERNETES_SERVICE_PORT_HTTPS=443
    KUBERNETES_PORT_443_TCP_PORT=443
    KUBERNETES_PORT_443_TCP=tcp://10.43.0.1:443
    #master节点
    MASTER_ADDR=pytorchjob-seg-1275-master-0
    #当前节点的RANK
    RANK=0
    PYTHONPATH=/app:
    KUBERNETES_SERVICE_PORT=443
    #由KubeFlow的training operator自动注入,通信的默认端口是23456
    MASTER_PORT=23456
    PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
    NCCL_DEBUG=INFO
    NCCL_SOCKET_IFNAME=eth0
    KUBERNETES_SERVICE_HOST=10.43.0.1
  • 对数据集设置DistributedSampler,并把DataLoader的sampler arg设置为创建的sampler

    sampler = torch.utils.data.distributed.DistributedSampler(dataset, rank=0, num_replicas=WORLD_SIZE)
    dataloader = DataLoader(dataset, args.batch_size, shuffle=False, sampler=sampler)
  • 对模型进行DDP封装,被DDP封装的model的参数的梯度才会进行all reduce

    model = nn.parallel.DistributedDataParallel(model)
  • 分布式下设置epoch

    sampler.set_epoch(epoch)

    示例代码:https://github.com/chengleqi/deeplabv3 ,仍可优化,欢迎提出PR改进。

Tensorflow