From f73a6ed635659f03ac583a1d914ea07a2cbeab99 Mon Sep 17 00:00:00 2001
From: Jiahao Li <liplus17@163.com>
Date: Tue, 17 Sep 2024 21:29:11 +0800
Subject: [PATCH] Enable overlap_comm for better performance (#846)

Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
---
 applications/DeepSpeed-Chat/dschat/utils/ds_utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/applications/DeepSpeed-Chat/dschat/utils/ds_utils.py b/applications/DeepSpeed-Chat/dschat/utils/ds_utils.py
index 9c15e5143..0cf1c28ab 100644
--- a/applications/DeepSpeed-Chat/dschat/utils/ds_utils.py
+++ b/applications/DeepSpeed-Chat/dschat/utils/ds_utils.py
@@ -33,6 +33,7 @@ def get_train_ds_config(offload,
         dtype_config = {"enabled": True}
     zero_opt_dict = {
         "stage": stage,
+        "overlap_comm": True,
         "offload_param": {
             "device": device
         },