[Eager] Support sharding_parallel under eager (#42910)

python/paddle/distributed/fleet/utils/hybrid_parallel_util.py

@@ -162,9 +162,15 @@ def sharding_reduce_gradients(parameter_list, hcg):
         sharding_nrank = hcg.get_sharding_parallel_group().nranks
         for param in parameter_list:
             if param.trainable and (param._grad_ivar() is not None):
+                if in_dygraph_mode():
+                    param.grad.scale_(1.0 / sharding_nrank)
+                    paddle.distributed.all_reduce(
+                        param.grad,
+                        group=hcg.get_sharding_parallel_group(),
+                        use_calc_stream=True)
 
+                elif _in_legacy_dygraph():
                     g_var = param._grad_ivar()
-
                     # need use trace_op to allreduce 
                     # paddle.distributed.all_reduce(
                     #     g_var, group=hcg.get_sharding_parallel_group(), use_calc_stream=True)
@@ -178,7 +184,8 @@ def sharding_reduce_gradients(parameter_list, hcg):
                         })
 
                     # grad / sharding_rank
-                div_factor = paddle.to_tensor(sharding_nrank, dtype=g_var.dtype)
+                    div_factor = paddle.to_tensor(
+                        sharding_nrank, dtype=g_var.dtype)
                     paddle.fluid.framework._dygraph_tracer().trace_op(
                         type="elementwise_div",
                         inputs={'X': g_var,

python/paddle/fluid/tests/unittests/test_parallel_dygraph_sharding_parallel.py

@@ -25,8 +25,7 @@ class TestHybridParallel(TestMultipleGpus):
 
     # check sharding logic as well as the accuracy with single mode
     def test_hybrid_parallel_sharding_logic(self):
-        # self.run_mnist_2gpu(
-        #     'hybrid_parallel_sharding_model.py')
+        self.run_mnist_2gpu('hybrid_parallel_sharding_model.py')
         self.run_mnist_2gpu(
             'hybrid_parallel_sharding_model.py', eager_mode=False)