set_deviceid before call gpu (#69633)

wanghuancoder · web-flow · commit a2395fec90d4 · 2024-11-25T14:04:24.000+08:00
diff --git a/paddle/fluid/pybind/eager_functions.cc b/paddle/fluid/pybind/eager_functions.cc
@@ -139,6 +139,7 @@ static PyObject* eager_api_scale(PyObject* self,
   paddle::Tensor ret;
   {
     eager_gil_scoped_release guard;
+    EagerSetDeviceId();
     ret = egr::scale(tensor, scale, bias, bias_after_scale, trace_backward);
   }
   return ToPyObject(ret);
@@ -159,6 +160,7 @@ static PyObject* eager_api_run_backward(PyObject* self,
   }
   {
     eager_gil_scoped_release guard;
+    EagerSetDeviceId();
     egr::Backward(tensors, grad_tensors, retain_graph);
   }
   RETURN_PY_NONE
@@ -189,6 +191,7 @@ static PyObject* eager_api_run_partial_grad(PyObject* self,
   std::vector<paddle::Tensor> result;
   {
     eager_gil_scoped_release guard;
+    EagerSetDeviceId();
     result = egr::Grad(tensors,
                        inputs,
                        grad_tensors,
@@ -216,6 +219,7 @@ static PyObject* eager_api_tensor_copy(PyObject* self,
 
   {
     eager_gil_scoped_release guard;
+    EagerSetDeviceId();
     dst = src.copy_to(place, blocking);
     egr::EagerUtils::autograd_meta(&dst)->SetStopGradient(
         egr::EagerUtils::autograd_meta(&(src))->StopGradient());
@@ -461,6 +465,7 @@ static PyObject* eager_api_jit_function_call(PyObject* self,
   std::vector<paddle::Tensor> outs;
   {
     eager_gil_scoped_release guard;
+    EagerSetDeviceId();
     outs = (*function)(ins);
   }
   return ToPyObject(outs);
@@ -672,6 +677,7 @@ PyObject* eager_api_run_custom_op(PyObject* self,
 
   {
     eager_gil_scoped_release guard;
+    EagerSetDeviceId();
     ctx.ConstructInplaceIndex(inputs, outputs, inplace_map);
     const auto& inplace_reverse_idx_map = ctx.GetInplaceReverseIndexMap();
     for (size_t out_idx = 0; out_idx < outputs.size(); ++out_idx) {
@@ -872,6 +878,7 @@ static PyObject* eager_api_sparse_coo_tensor(PyObject* self,
   paddle::Tensor tensor;
   {
     eager_gil_scoped_release guard;
+    EagerSetDeviceId();
     PADDLE_ENFORCE(
         non_zero_indices.is_dense_tensor(),
         common::errors::Fatal("the non-zero indices must be a DenseTensor."));
@@ -916,6 +923,7 @@ static PyObject* eager_api_sparse_csr_tensor(PyObject* self,
   paddle::Tensor tensor;
   {
     eager_gil_scoped_release guard;
+    EagerSetDeviceId();
     PADDLE_ENFORCE(non_zero_crows.is_dense_tensor(),
                    common::errors::Fatal(
                        "the compressed non-zero rows must be a DenseTensor."));
@@ -996,6 +1004,7 @@ static PyObject* eager_api_async_read(PyObject* self,
 
   {
     eager_gil_scoped_release guard;
+    EagerSetDeviceId();
     PADDLE_ENFORCE_EQ(
         src.is_gpu_pinned(),
         true,
@@ -1175,6 +1184,7 @@ static PyObject* eager_api_async_write(PyObject* self,
   }
   {
     eager_gil_scoped_release guard;
+    EagerSetDeviceId();
     PADDLE_ENFORCE_EQ(
         src.is_gpu(),
         true,
diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
@@ -619,6 +619,9 @@ static PyObject* tensor_method__copy_to(TensorObject* self,
   paddle::Tensor cp_tensor;
   {
     eager_gil_scoped_release guard;
+
+    EagerSetDeviceId();
+
     cp_tensor = self->tensor.copy_to(place, blocking);
     if (!blocking) {
       IncreaseTensorReferenceCountUntilCopyComplete(self->tensor, place);
@@ -690,6 +693,9 @@ static PyObject* tensor_method_copy_(TensorObject* self,
           << self->tensor.name();
   if (!self->tensor.initialized()) {
     eager_gil_scoped_release guard;
+
+    EagerSetDeviceId();
+
     egr::EagerUtils::autograd_meta(&(self->tensor))
         ->SetStopGradient(
             egr::EagerUtils::autograd_meta(&(src_tensor))->StopGradient());
@@ -702,6 +708,9 @@ static PyObject* tensor_method_copy_(TensorObject* self,
   } else {
     if (src_tensor.has_allocation()) {
       eager_gil_scoped_release guard;
+
+      EagerSetDeviceId();
+
       self->tensor.copy_(src_tensor, self->tensor.place(), blocking);
     }
   }
@@ -766,6 +775,9 @@ static PyObject* tensor_method_clone(TensorObject* self,
   paddle::Tensor out;
   {
     eager_gil_scoped_release guard;
+
+    EagerSetDeviceId();
+
     PADDLE_ENFORCE_EQ(
         self->tensor.initialized(),
         true,
@@ -922,6 +934,7 @@ static PyObject* tensor_clear_gradient(TensorObject* self,
                   ->unsafe_mutable_value();
         }
         if (set_to_zero) {
+          EagerSetDeviceId();
           auto* dev_ctx =
               phi::DeviceContextPool::Instance().Get(grad_t->place());
           phi::funcs::set_constant(*dev_ctx, grad_t, 0.0);
@@ -952,6 +965,7 @@ static PyObject* tensor__zero_grads(TensorObject* self,
 
   if (egr::EagerUtils::IsLeafTensor(self->tensor)) {
     eager_gil_scoped_release guard;
+    EagerSetDeviceId();
     // Add RetainGrad as PostHook to AccumulationNode
     paddle::Tensor* grad = egr::EagerUtils::mutable_grad(self->tensor);
     PADDLE_ENFORCE(
@@ -976,6 +990,7 @@ static PyObject* tensor__zero_grads(TensorObject* self,
     }
   } else {
     eager_gil_scoped_release guard;
+    EagerSetDeviceId();
     auto meta = egr::EagerUtils::unsafe_autograd_meta(self->tensor);
     if (meta->MutableGrad()->initialized()) {
       if (meta->MutableGrad()->is_dense_tensor() ||
@@ -3250,6 +3265,7 @@ static PyObject* tensor_contiguous(TensorObject* self,
       return reinterpret_cast<PyObject*>(self);
     } else {
       eager_gil_scoped_release guard;
+      EagerSetDeviceId();
       *dense_tensor = paddle::experimental::Trans2Contiguous(*dense_tensor);
       Py_INCREF(self);
       return reinterpret_cast<PyObject*>(self);
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -2797,4 +2797,39 @@ CvtPlacements(Placements placements, int ndim) {
   return {dim_map, partial_status};
 }
 
+void EagerSetDeviceId() {
+  auto expected_place = egr::Controller::Instance().GetExpectedPlace();
+
+  if (phi::is_gpu_place(expected_place)) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+    phi::backends::gpu::SetDeviceId(expected_place.device);
+    VLOG(4) << "CurrentDeviceId: " << phi::backends::gpu::GetCurrentDeviceId()
+            << " from " << (int)expected_place.device;  // NOLINT
+#else
+    PADDLE_THROW(common::errors::PreconditionNotMet(
+        "PaddlePaddle should compile with GPU if use CUDAPlace."));
+#endif
+  } else if (phi::is_custom_place(expected_place)) {
+#if defined(PADDLE_WITH_CUSTOM_DEVICE)
+    phi::DeviceManager::SetDevice(expected_place);
+    VLOG(4) << "CurrentDeviceId: "
+            << phi::DeviceManager::GetDevice(expected_place.GetDeviceType())
+            << " from " << (int)expected_place.device;  // NOLINT
+#else
+    PADDLE_THROW(common::errors::PreconditionNotMet(
+        "PaddlePaddle should compile with CUSTOM_DEVICE if use CustomPlace."));
+#endif
+  } else if (phi::is_xpu_place(expected_place)) {
+#if defined(PADDLE_WITH_XPU)
+    phi::backends::xpu::SetXPUDeviceId(expected_place.device);
+    VLOG(4) << "CurrentDeviceId: "
+            << phi::backends::xpu::GetXPUCurrentDeviceId() << " from "
+            << (int)expected_place.device;  // NOLINT
+#else
+    PADDLE_THROW(common::errors::PreconditionNotMet(
+        "PaddlePaddle should compile with XPU if use XPUPlace."));
+#endif
+  }
+}
+
 }  // namespace paddle::pybind
diff --git a/paddle/fluid/pybind/eager_utils.h b/paddle/fluid/pybind/eager_utils.h
@@ -449,5 +449,7 @@ std::tuple<std::vector<int64_t>,
            paddle::flat_hash_map<int64_t, phi::ReduceType>>
 CvtPlacements(phi::distributed::Placements placements, int ndim);
 
+void EagerSetDeviceId();
+
 }  // namespace pybind
 }  // namespace paddle

Original file line number	Diff line number	Diff line change
`@@ -139,6 +139,7 @@ static PyObject* eager_api_scale(PyObject* self,`
`139`	`139`	`paddle::Tensor ret;`
`140`	`140`	`{`
`141`	`141`	`eager_gil_scoped_release guard;`
	`142`	`+ EagerSetDeviceId();`
`142`	`143`	`ret = egr::scale(tensor, scale, bias, bias_after_scale, trace_backward);`
`143`	`144`	`}`
`144`	`145`	`return ToPyObject(ret);`
`@@ -159,6 +160,7 @@ static PyObject* eager_api_run_backward(PyObject* self,`
`159`	`160`	`}`
`160`	`161`	`{`
`161`	`162`	`eager_gil_scoped_release guard;`
	`163`	`+ EagerSetDeviceId();`
`162`	`164`	`egr::Backward(tensors, grad_tensors, retain_graph);`
`163`	`165`	`}`
`164`	`166`	`RETURN_PY_NONE`
`@@ -189,6 +191,7 @@ static PyObject* eager_api_run_partial_grad(PyObject* self,`
`189`	`191`	`std::vector<paddle::Tensor> result;`
`190`	`192`	`{`
`191`	`193`	`eager_gil_scoped_release guard;`
	`194`	`+ EagerSetDeviceId();`
`192`	`195`	`result = egr::Grad(tensors,`
`193`	`196`	`inputs,`
`194`	`197`	`grad_tensors,`
`@@ -216,6 +219,7 @@ static PyObject* eager_api_tensor_copy(PyObject* self,`
`216`	`219`
`217`	`220`	`{`
`218`	`221`	`eager_gil_scoped_release guard;`
	`222`	`+ EagerSetDeviceId();`
`219`	`223`	`dst = src.copy_to(place, blocking);`
`220`	`224`	`egr::EagerUtils::autograd_meta(&dst)->SetStopGradient(`
`221`	`225`	`egr::EagerUtils::autograd_meta(&(src))->StopGradient());`
`@@ -461,6 +465,7 @@ static PyObject* eager_api_jit_function_call(PyObject* self,`
`461`	`465`	`std::vector<paddle::Tensor> outs;`
`462`	`466`	`{`
`463`	`467`	`eager_gil_scoped_release guard;`
	`468`	`+ EagerSetDeviceId();`
`464`	`469`	`outs = (*function)(ins);`
`465`	`470`	`}`
`466`	`471`	`return ToPyObject(outs);`
`@@ -672,6 +677,7 @@ PyObject* eager_api_run_custom_op(PyObject* self,`
`672`	`677`
`673`	`678`	`{`
`674`	`679`	`eager_gil_scoped_release guard;`
	`680`	`+ EagerSetDeviceId();`
`675`	`681`	`ctx.ConstructInplaceIndex(inputs, outputs, inplace_map);`
`676`	`682`	`const auto& inplace_reverse_idx_map = ctx.GetInplaceReverseIndexMap();`
`677`	`683`	`for (size_t out_idx = 0; out_idx < outputs.size(); ++out_idx) {`
`@@ -872,6 +878,7 @@ static PyObject* eager_api_sparse_coo_tensor(PyObject* self,`
`872`	`878`	`paddle::Tensor tensor;`
`873`	`879`	`{`
`874`	`880`	`eager_gil_scoped_release guard;`
	`881`	`+ EagerSetDeviceId();`
`875`	`882`	`PADDLE_ENFORCE(`
`876`	`883`	`non_zero_indices.is_dense_tensor(),`
`877`	`884`	`common::errors::Fatal("the non-zero indices must be a DenseTensor."));`
`@@ -916,6 +923,7 @@ static PyObject* eager_api_sparse_csr_tensor(PyObject* self,`
`916`	`923`	`paddle::Tensor tensor;`
`917`	`924`	`{`
`918`	`925`	`eager_gil_scoped_release guard;`
	`926`	`+ EagerSetDeviceId();`
`919`	`927`	`PADDLE_ENFORCE(non_zero_crows.is_dense_tensor(),`
`920`	`928`	`common::errors::Fatal(`
`921`	`929`	`"the compressed non-zero rows must be a DenseTensor."));`
`@@ -996,6 +1004,7 @@ static PyObject* eager_api_async_read(PyObject* self,`
`996`	`1004`
`997`	`1005`	`{`
`998`	`1006`	`eager_gil_scoped_release guard;`
	`1007`	`+ EagerSetDeviceId();`
`999`	`1008`	`PADDLE_ENFORCE_EQ(`
`1000`	`1009`	`src.is_gpu_pinned(),`
`1001`	`1010`	`true,`
`@@ -1175,6 +1184,7 @@ static PyObject* eager_api_async_write(PyObject* self,`
`1175`	`1184`	`}`
`1176`	`1185`	`{`
`1177`	`1186`	`eager_gil_scoped_release guard;`
	`1187`	`+ EagerSetDeviceId();`
`1178`	`1188`	`PADDLE_ENFORCE_EQ(`
`1179`	`1189`	`src.is_gpu(),`
`1180`	`1190`	`true,`