From 6b57cc705f07ac5ea5e18fffe76b4470669b1039 Mon Sep 17 00:00:00 2001 From: laiyongqiang Date: Thu, 5 Nov 2020 20:02:01 +0800 Subject: [PATCH] disable somas when variable_memory_max_size is set --- .../ccsrc/runtime/device/kernel_runtime.cc | 17 ++++++++++------- .../official/nlp/bert_thor/run_pretrain.py | 1 - .../nlp/tinybert/run_general_distill.py | 1 - 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/mindspore/ccsrc/runtime/device/kernel_runtime.cc b/mindspore/ccsrc/runtime/device/kernel_runtime.cc index 61f9682e38..64f1da2fb3 100644 --- a/mindspore/ccsrc/runtime/device/kernel_runtime.cc +++ b/mindspore/ccsrc/runtime/device/kernel_runtime.cc @@ -675,13 +675,16 @@ void KernelRuntime::AssignDynamicMemory(session::KernelGraph *graph) { if (is_enable_mem_reuse) { MS_LOG(INFO) << "Memory Reuse is enable..."; -#ifdef MEM_REUSE_DEBUG - mem_manager_->MallocReusedDynamicMem(graph); - mem_type = kReuseDynamicMem; -#else - mem_manager_->MallocSomasDynamicMem(graph); - mem_type = kSomasReuseDynamicMem; -#endif + auto context = MsContext::GetInstance(); + MS_EXCEPTION_IF_NULL(context); + auto variable_memory_max_size = context->get_param(MS_CTX_VARIABLE_MEMORY_MAX_SIZE); + if (variable_memory_max_size == "0") { + mem_manager_->MallocSomasDynamicMem(graph); + mem_type = kSomasReuseDynamicMem; + } else { + mem_manager_->MallocReusedDynamicMem(graph); + mem_type = kReuseDynamicMem; + } } else { MS_LOG(INFO) << "Memory Reuse is disable..."; } diff --git a/model_zoo/official/nlp/bert_thor/run_pretrain.py b/model_zoo/official/nlp/bert_thor/run_pretrain.py index e00eb2de79..410878e08a 100644 --- a/model_zoo/official/nlp/bert_thor/run_pretrain.py +++ b/model_zoo/official/nlp/bert_thor/run_pretrain.py @@ -145,7 +145,6 @@ def run_pretrain(): context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target, device_id=args_opt.device_id, save_graphs=False) context.set_context(reserve_class_name_in_scope=False) - context.set_context(variable_memory_max_size="30GB") context.set_context(max_call_depth=3000) ckpt_save_dir = args_opt.save_checkpoint_path if args_opt.distribute == "true": diff --git a/model_zoo/official/nlp/tinybert/run_general_distill.py b/model_zoo/official/nlp/tinybert/run_general_distill.py index ceb9b82ae7..2bfb3f417f 100644 --- a/model_zoo/official/nlp/tinybert/run_general_distill.py +++ b/model_zoo/official/nlp/tinybert/run_general_distill.py @@ -70,7 +70,6 @@ def run_general_distill(): raise Exception("Target error, GPU or Ascend is supported.") context.set_context(reserve_class_name_in_scope=False) - context.set_context(variable_memory_max_size="30GB") save_ckpt_dir = os.path.join(args_opt.save_ckpt_path, datetime.datetime.now().strftime('%Y-%m-%d_time_%H_%M_%S'))