From 60fc06c610da7aada4893fc55b8871ed7ed05c94 Mon Sep 17 00:00:00 2001
From: Reza Yazdani <44502768+RezaYazdaniAminabadi@users.noreply.github.com>
Date: Thu, 3 Mar 2022 14:56:30 -0800
Subject: [PATCH] Synchronize the GPUs for the text-generation inference test
 (#1805)

---
 deepspeed/inference/engine.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py
index 156812c1..3cd24894 100755
--- a/deepspeed/inference/engine.py
+++ b/deepspeed/inference/engine.py
@@ -356,6 +356,7 @@ class InferenceEngine(Module):
                         input = input.to(torch.cuda.current_device())
                         if not input.is_contiguous():
                             input = input.contiguous()
+                        dist.broadcast(input, 0)
                 for k in kwargs:
                     if torch.is_tensor(kwargs[k]):
                         kwargs[k] = kwargs[k].to(torch.cuda.current_device())
-- 
GitLab