From 207eddd73cfff8295bd70d315ee9b66d19b10e1b Mon Sep 17 00:00:00 2001
From: lizhenyu <lizhenyu13@huawei.com>
Date: Mon, 2 Aug 2021 16:29:47 +0800
Subject: [PATCH] [bugfix]sync stream failed in ps cache training on Ascend

---
 mindspore/ccsrc/ps/ps_cache/ps_cache_manager.cc | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/mindspore/ccsrc/ps/ps_cache/ps_cache_manager.cc b/mindspore/ccsrc/ps/ps_cache/ps_cache_manager.cc
index 4c4e97f7939..fa033043716 100644
--- a/mindspore/ccsrc/ps/ps_cache/ps_cache_manager.cc
+++ b/mindspore/ccsrc/ps/ps_cache/ps_cache_manager.cc
@@ -237,9 +237,6 @@ void PsCacheManager::AllocMemForHashTable() {
   embedding_device_cache_->hash_swap_value_addr_ = reinterpret_cast<float *>(
     embedding_device_cache_->cache_->MallocMemory(max_embedding_size * batch_elements_ * sizeof(float)));
   MS_EXCEPTION_IF_NULL(embedding_device_cache_->hash_swap_value_addr_);
-  if (!(embedding_device_cache_->cache_->MallocConstantMemory(vocab_cache_size_))) {
-    MS_LOG(EXCEPTION) << "MallocConstantMemory failed.";
-  }
 }
 
 void PsCacheManager::SetLocalIdRank() {
@@ -328,6 +325,14 @@ void PsCacheManager::ProcessDataTask(uint32_t device_id, const void *context) {
   MS_ERROR_IF_NULL_WO_RET_VAL(embedding_device_cache_);
   MS_ERROR_IF_NULL_WO_RET_VAL(embedding_device_cache_->cache_);
   embedding_device_cache_->cache_->InitDevice(device_id, context);
+
+  // MallocConstantMemory need stream on device Ascend, should be called after InitDevice.
+  if (!(embedding_device_cache_->cache_->MallocConstantMemory(vocab_cache_size_))) {
+    MS_LOG(ERROR) << "MallocConstantMemory failed.";
+    running_ = false;
+    return;
+  }
+
   InitParameterServer();
   InitDataChannel();
   while (running_) {