Как совместить общение NCCL с обучением PyTorch

Question

Я использую PyTorch для распределенного обучения модели. Что я хочу сделать, так это наложить связь NCCL на обратную (вычислительную) модель. Я использую многопоточность python. Однако кажется, что когда два потока работают вместе, время обратного движения увеличивается.

Вот часть моего кода, связанная с этой проблемой. Я упростил свой код, чтобы он был понятнее.

semaphore = threading.Semaphore(0)

def bk_hook(*args):
    semaphore.release()

for layer in model.__dict__['_modules'].items():
    layer[1].register_backward_hook(bk_hook)

def training():
    # trains the model

def communication():
    global model, semaphore
    layers = len(list(model.__dict__['_modules']))
    for layer in range(layers):
        semaphore.acquire()
        dist.broadcast(...) # NCCL communication

t1 = threading.Thread(target=training)
t2 = threading.Thread(target=communication)
t1.start()
t2.start()

Я пробовал запустить поток, который только обучает модель, и другой поток, который не использует NCCL. В этом случае время обратного хода не увеличивается. Таким образом, именно коммуникация NCCL приводит к увеличению обратного времени. Я использую 3 узла на одном сервере, и в этом случае NCCL использует cudaMemcpy. Интересно, почему коммуникация и вычисления не могут пересекаться.