Make multiprocessing_train.py work with multi-node setups

# can be found in the PATENTS file in the same directory.

import os

import random

import signal

import torch

def main(args):

    # Set distributed training parameters for a single node.

    args.distributed_world_size = torch.cuda.device_count()

    port = random.randint(10000, 20000)

    args.distributed_init_method = 'tcp://localhost:{port}'.format(port=port)

    args.distributed_init_host = 'localhost'

    args.distributed_port = port + 1

    if max(args.update_freq) > 1 and args.ddp_backend != 'no_c10d':

        print('| WARNING: when using --update-freq on a single machine, you '

              'will get better performance with --ddp-backend=no_c10d')

    # Train with multiprocessing.

    procs = []

    for i in range(args.distributed_world_size):

        args.distributed_rank = i

    base_rank = args.distributed_rank

    for i in range(torch.cuda.device_count()):

        args.distributed_rank = base_rank + i

        args.device_id = i

        procs.append(mp.Process(target=run, args=(args, error_queue, ), daemon=True))

        procs[i].start()

import itertools

import os

import math

import random

import torch

from fairseq import distributed_utils, options, progress_bar, tasks, utils

    elif args.distributed_world_size > 1:

        from multiprocessing_train import main as multiprocessing_main

        # Set distributed training parameters for a single node.

        args.distributed_world_size = torch.cuda.device_count()

        port = random.randint(10000, 20000)

        args.distributed_init_method = 'tcp://localhost:{port}'.format(port=port)

        args.distributed_port = port + 1

        multiprocessing_main(args)

    else:

        main(args)