diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc index 0f539cc6e9..2f0c0886f1 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc @@ -1092,6 +1092,25 @@ void PynativeExecutor::CleanPreMemoryInValueNode(const std::string &cell_id) { tensor->set_device_address(nullptr); } } + + if (dynamic_cell_) { + std::set forward_op_tensor_id; + for (const auto &elem : cell_op_index_with_tensor_id_[top_cell_id_]) { + const auto &tensor_id_list = elem.second; + for (const auto &tensor_id : tensor_id_list) { + forward_op_tensor_id.emplace(tensor_id); + } + } + for (auto &tensor : all_value_node_tensors_) { + if (tensor->device_address() != nullptr && + forward_op_tensor_id.find(tensor->id()) != forward_op_tensor_id.end()) { + tensor->device_address()->ClearDeviceMemory(); + tensor->set_device_address(nullptr); + } + } + all_value_node_tensors_.clear(); + } + top_cell_id_ = cell_id; } @@ -2086,6 +2105,37 @@ std::string PynativeExecutor::GetGradCellId(bool has_sens, const py::object &cel return cell_id; } +void PynativeExecutor::SaveAllValueNodeTensors(const FuncGraphPtr &graph) { + std::unordered_set all_value_node_tensors; + auto trace_function = [&all_value_node_tensors](const AnfNodePtr &anf_node) { + auto value = GetValueNode(anf_node); + if (value) { + if (value->isa()) { + auto tensor = value->cast(); + MS_EXCEPTION_IF_NULL(tensor); + if (tensor->device_address()) { + all_value_node_tensors.emplace(tensor); + } + } else if (value->isa()) { + auto tuple = value->cast(); + MS_EXCEPTION_IF_NULL(tuple); + for (size_t i = 0; i < tuple->size(); i++) { + if ((*tuple)[i]->isa()) { + auto tensor = (*tuple)[i]->cast(); + MS_EXCEPTION_IF_NULL(tensor); + if (tensor->device_address()) { + all_value_node_tensors.emplace(tensor); + } + } + } + } + } + return FOLLOW; + }; + (void)TopoSort(graph->get_return(), SuccDeeperSimple, trace_function); + all_value_node_tensors_ = all_value_node_tensors; +} + void PynativeExecutor::GradNetInner(const GradOperationPtr &grad, const py::object &cell, const py::object &weights, const py::args &args) { auto size = args.size(); @@ -2127,6 +2177,9 @@ void PynativeExecutor::GradNetInner(const GradOperationPtr &grad, const py::obje resource->results()[pipeline::kBackend] = compile::CreateBackend(); MS_LOG(INFO) << "Start opt"; + if (dynamic_cell_) { + SaveAllValueNodeTensors(resource->func_graph()); + } PynativeOptimizeAction(resource); SaveTensorsInValueNode(resource); TaskEmitAction(resource); diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h index b74a851096..b33e516337 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h @@ -198,6 +198,7 @@ class PynativeExecutor : public std::enable_shared_from_this { // Update the abstract and device address info of value node and tensors in bprop graph void UpdateAbstractAndDeviceAddress(const OpExecInfoPtr &op_exec_info, const py::object &out_real); void SaveTensorsInValueNode(const ResourcePtr &resource); + void SaveAllValueNodeTensors(const FuncGraphPtr &graph); void CleanPreMemoryInValueNode(const std::string &cell_id); // Construct grad graph @@ -304,6 +305,7 @@ class PynativeExecutor : public std::enable_shared_from_this { std::unordered_map cell_tensor_id_with_tensor_; std::unordered_map node_abs_map_; std::unordered_map prim_abs_list_; + std::unordered_set all_value_node_tensors_; }; using PynativeExecutorPtr = std::shared_ptr; diff --git a/mindspore/ccsrc/runtime/device/ascend/ascend_device_address.cc b/mindspore/ccsrc/runtime/device/ascend/ascend_device_address.cc index b00a31528d..60f1c39fe6 100644 --- a/mindspore/ccsrc/runtime/device/ascend/ascend_device_address.cc +++ b/mindspore/ccsrc/runtime/device/ascend/ascend_device_address.cc @@ -612,7 +612,7 @@ bool AscendDeviceAddress::ConvertFormatAndSyncHostToDevice(const ShapeVector &sh return sync_ok; } -AscendDeviceAddress::~AscendDeviceAddress() { +void AscendDeviceAddress::ClearDeviceMemory() { if (ptr_ == nullptr) { return; } @@ -627,6 +627,8 @@ AscendDeviceAddress::~AscendDeviceAddress() { } } +AscendDeviceAddress::~AscendDeviceAddress() { ClearDeviceMemory(); } + bool AscendDeviceAddress::DumpMemToFile(bool trans_flag, const std::string &filepath, const std::string &host_fmt, const ShapeVector &host_shape, TypeId host_type) const { bool ret = false; diff --git a/mindspore/ccsrc/runtime/device/ascend/ascend_device_address.h b/mindspore/ccsrc/runtime/device/ascend/ascend_device_address.h index 393525c2e3..819e1a325a 100644 --- a/mindspore/ccsrc/runtime/device/ascend/ascend_device_address.h +++ b/mindspore/ccsrc/runtime/device/ascend/ascend_device_address.h @@ -41,6 +41,7 @@ class AscendDeviceAddress : public DeviceAddress { ~AscendDeviceAddress() override; bool SyncDeviceToHost(const ShapeVector &shape, size_t size, TypeId type, void *host_ptr) const override; bool SyncHostToDevice(const ShapeVector &shape, size_t size, TypeId type, const void *host_ptr) const override; + void ClearDeviceMemory() override; DeviceAddressType DeviceType() const override { return DeviceAddressType::kAscend; } bool DumpMemToFile(bool dump_mode, const std::string &filepath, const std::string &host_fmt, const ShapeVector &host_shape, TypeId host_type) const override; diff --git a/mindspore/ccsrc/runtime/device/cpu/cpu_device_address.h b/mindspore/ccsrc/runtime/device/cpu/cpu_device_address.h index d73804c324..bf8230ec35 100644 --- a/mindspore/ccsrc/runtime/device/cpu/cpu_device_address.h +++ b/mindspore/ccsrc/runtime/device/cpu/cpu_device_address.h @@ -35,6 +35,7 @@ class CPUDeviceAddress : public DeviceAddress { bool SyncDeviceToHost(const ShapeVector &shape, size_t size, TypeId type, void *host_ptr) const override; bool SyncHostToDevice(const ShapeVector &shape, size_t size, TypeId type, const void *host_ptr) const override; + void ClearDeviceMemory() override {} DeviceAddressType DeviceType() const override { return DeviceAddressType::kCPU; } }; } // namespace cpu diff --git a/mindspore/ccsrc/runtime/device/gpu/gpu_device_address.cc b/mindspore/ccsrc/runtime/device/gpu/gpu_device_address.cc index af2570598b..4893ebdc38 100644 --- a/mindspore/ccsrc/runtime/device/gpu/gpu_device_address.cc +++ b/mindspore/ccsrc/runtime/device/gpu/gpu_device_address.cc @@ -69,7 +69,7 @@ bool GPUDeviceAddress::SyncHostToDevice(const ShapeVector &, size_t size, TypeId return GPUDeviceManager::GetInstance().SyncStream(stream); } -GPUDeviceAddress::~GPUDeviceAddress() { +void GPUDeviceAddress::ClearDeviceMemory() { if (ptr_ == nullptr) { return; } @@ -78,6 +78,8 @@ GPUDeviceAddress::~GPUDeviceAddress() { ptr_ = nullptr; } } + +GPUDeviceAddress::~GPUDeviceAddress() { ClearDeviceMemory(); } #ifdef ENABLE_DEBUGGER bool GPUDeviceAddress::LoadMemToHost(const std::string &tensor_name, int execution_order, const std::string &host_fmt, const ShapeVector &host_shape, TypeId host_type, size_t slot, diff --git a/mindspore/ccsrc/runtime/device/gpu/gpu_device_address.h b/mindspore/ccsrc/runtime/device/gpu/gpu_device_address.h index a98f67786b..943ca8e596 100644 --- a/mindspore/ccsrc/runtime/device/gpu/gpu_device_address.h +++ b/mindspore/ccsrc/runtime/device/gpu/gpu_device_address.h @@ -38,6 +38,7 @@ class GPUDeviceAddress : public DeviceAddress { bool SyncDeviceToHost(const ShapeVector &shape, size_t size, TypeId type, void *host_ptr) const override; bool SyncHostToDevice(const ShapeVector &shape, size_t size, TypeId type, const void *host_ptr) const override; + void ClearDeviceMemory() override; void set_status(DeviceAddressStatus status) { status_ = status; } DeviceAddressStatus status() const { return status_; } DeviceAddressType DeviceType() const override { return DeviceAddressType::kGPU; } diff --git a/mindspore/core/ir/device_sync.h b/mindspore/core/ir/device_sync.h index 2cf7ecd38e..766049370e 100644 --- a/mindspore/core/ir/device_sync.h +++ b/mindspore/core/ir/device_sync.h @@ -33,6 +33,7 @@ class DeviceSync { virtual bool SyncDeviceToHost(const ShapeVector &shape, size_t size, TypeId type, void *host_ptr) const = 0; virtual bool SyncHostToDevice(const ShapeVector &shape, size_t size, TypeId type, const void *host_ptr) const = 0; virtual void *GetMutablePtr() const = 0; + virtual void ClearDeviceMemory() = 0; }; using DeviceSyncPtr = std::shared_ptr; } // namespace mindspore