[CustomDevice] fix release ProcessGroupCustom error when process exit (#59800)

ronny1996 · web-flow · commit bb66cb1b0f31 · 2023-12-08T17:40:33.000+08:00
diff --git a/paddle/fluid/memory/allocation/custom_allocator.cc b/paddle/fluid/memory/allocation/custom_allocator.cc
@@ -29,8 +29,10 @@ void CustomAllocator::FreeImpl(phi::Allocation* allocation) {
       platform::errors::PermissionDenied("CustomDevice memory is "
                                          "freed in incorrect device. "
                                          "This may be a bug"));
-  phi::DeviceManager::GetDeviceWithPlace(place_)->MemoryDeallocate(
-      allocation->ptr(), allocation->size());
+  if (phi::DeviceManager::HasDeviceType(place_.GetDeviceType())) {
+    phi::DeviceManager::GetDeviceWithPlace(place_)->MemoryDeallocate(
+        allocation->ptr(), allocation->size());
+  }
   delete allocation;
 }
 
diff --git a/paddle/fluid/pybind/distributed_py.cc b/paddle/fluid/pybind/distributed_py.cc
@@ -1271,6 +1271,7 @@ void BindDistributed(py::module *m) {
                   py::arg("rank"),
                   py::arg("world_size"),
                   py::arg("group_id") = 0,
+                  py::return_value_policy::reference_internal,
                   py::call_guard<py::gil_scoped_release>());
 
 #endif
