[AutoParallel] Fix reshard_grad with partial (#73236)

Xing-lil · web-flow · commit ac7aa197f3e9 · 2025-06-12T10:59:35.000+08:00
diff --git a/paddle/fluid/eager/api/manual/eager_manual/nodes/reshard_node.cc b/paddle/fluid/eager/api/manual/eager_manual/nodes/reshard_node.cc
@@ -46,9 +46,10 @@ ReshardGradNode::operator()(
 
   // Collect GradIn Tensors, Attrs and Recovered TensorWrappers
   auto input = egr::EagerUtils::RecoverTensorWrapper(&this->input_);
-  const auto& dist_attr =
+  auto dist_attr =
       std::static_pointer_cast<phi::distributed::DistTensor>(input.impl())
           ->dist_attr();
+  dist_attr.clean_partial_status();
   auto& grad_out = hooked_grad[0][0];
   // Prepare Grad function call
 
diff --git a/test/auto_parallel/reshard_api.py b/test/auto_parallel/reshard_api.py
@@ -37,6 +37,8 @@ def run_test_cases(self):
         self.test_case_p_to_r()
         self.test_case_r_to_s()
         self.test_case_forward_and_backward()
+        self.test_case_p_to_s_reshard_grad()
+        self.test_case_p_to_r_reshard_grad()
 
     def test_case_p_to_r(self):
         a = paddle.ones(self._shape)
@@ -126,6 +128,64 @@ def test_case_forward_and_backward(self):
             atol=1e-5,
         )
 
+    def test_case_p_to_s_reshard_grad(self):
+        if self._backend == "cpu":
+            return
+
+        np.random.seed(1901)
+        input_numpy = np.random.random(self._shape).astype("float32")
+        label_numpy = np.random.random(self._shape).astype('float32')
+
+        dist_input = dist.shard_tensor(
+            paddle.to_tensor(input_numpy),
+            dist.ProcessMesh([0, 1], dim_names=["x"]),
+            [Partial()],
+        )
+        dist_input.stop_gradient = False
+        dist_output = dist.reshard(
+            dist_input, dist.ProcessMesh([0, 1], dim_names=["x"]), [Shard(0)]
+        )
+        dist_label = dist.shard_tensor(
+            paddle.to_tensor(label_numpy),
+            dist.ProcessMesh([0, 1], dim_names=["x"]),
+            [Shard(0)],
+        )
+
+        dist_loss_fn = nn.MSELoss()
+        dist_loss = dist_loss_fn(dist_output, dist_label)
+        dist_loss.backward()
+
+        np.testing.assert_equal(dist_input.grad.placements, [dist.Replicate()])
+
+    def test_case_p_to_r_reshard_grad(self):
+        if self._backend == "cpu":
+            return
+
+        np.random.seed(1901)
+        input_numpy = np.random.random(self._shape).astype("float32")
+        label_numpy = np.random.random(self._shape).astype('float32')
+
+        dist_input = dist.shard_tensor(
+            paddle.to_tensor(input_numpy),
+            dist.ProcessMesh([0, 1], dim_names=["x"]),
+            [Partial()],
+        )
+        dist_input.stop_gradient = False
+        dist_output = dist.reshard(
+            dist_input, dist.ProcessMesh([0, 1], dim_names=["x"]), [Replicate()]
+        )
+        dist_label = dist.shard_tensor(
+            paddle.to_tensor(label_numpy),
+            dist.ProcessMesh([0, 1], dim_names=["x"]),
+            [Replicate()],
+        )
+
+        dist_loss_fn = nn.MSELoss()
+        dist_loss = dist_loss_fn(dist_output, dist_label)
+        dist_loss.backward()
+
+        np.testing.assert_equal(dist_input.grad.placements, [dist.Replicate()])
+
 
 if __name__ == '__main__':
     TestReshardAPI().run_test_cases()
