Использование нескольких графических процессоров для совместного обучения двух моделей

Question

В моей программе мне нужно совместно обучать две разные модели, которые имеют встраивание слова.Я должен использовать мульти-GPU для обучения, потому что мои данные слишком много.Если бы он тренировался в одном GPU, он провел бы более 7 дней.Поэтому моя идея описана ниже:

введите описание изображения здесь

mdoel-0 распределяется между GPU-0 и GPU-1 и моделью-1 работает одинаково на GPU-2 и GPU-3.Кроме того, обе модели имеют общие параметры встраивания слов, которые будут синхронно обновляться ими.

Я написал черновик об этом следующим образом.В этом проекте используются преимущества многопроцессорной обработки и распределенного обучения для нескольких графических процессоров на одном компьютере.Функции run_0 и run_1 относятся к модели-0 и модели-1 соответственно.

import torch
args.distributed_word_size=2
mp = torch.multiprocessing.get_context('spawn')
args.distributed_init_method_run_0 = 'tcp://localhost:00000'
args.distributed_init_method_run_1 = 'tcp://localhost:11111'

procs = []
for i in range(args.distributed_word_size):
    args.distributed_rank = i
    args.device_id = i
    procs.append(mp.Process(target=run_0, args=(args, ), daemon=True))
    procs[i].start()
for i in range(args.distributed_word_size):
    args.distributed_rank = i
    args.device_id = i + args.distributed_word_size
    procs.append(mp.Process(target=run_1, args=(args, ), daemon=True))
    procs[i+args.distributed_word_size].start()
for p in procs:
    p.joint()

def run_0(args):
    torch.distributed.init_process_group(
            backend=args.distributed_backend,
            init_method=args.distributed_init_method_run_0,
            world_size=args.distributed_world_size,
            rank=args.distributed_rank)
    main_0(args)
    ....

def run_1(args):
    torch.distributed.init_process_group(
            backend=args.distributed_backend,
            init_method=args.distributed_init_method_run_1,
            world_size=args.distributed_world_size,
            rank=args.distributed_rank)
    main_1(args)
    ....

Интересно, верен ли приведенный выше фрагмент кода и есть ли лучшие решения?В частности, как реализовать вложение общих слов для синхронного обновления в двух моделях?