[develop] Support release_grads in Pipeline Parallel and Sharding stage1 v1/v2 (#59739)

haohongxiang · web-flow · commit 9b40c67f672f · 2023-12-26T10:32:43.000+08:00
diff --git a/paddle/fluid/framework/distributed_strategy.proto b/paddle/fluid/framework/distributed_strategy.proto
@@ -81,6 +81,7 @@ message PpConfig {
     optional bool enable_timer = 3 [ default = false ];
     optional bool sharding_comm_overlap = 4 [ default = false ];
     optional bool profiling = 5 [ default = false ];
+    optional bool release_gradients = 6 [ default = false ];
 }
 
 message DygraphShardingConfig {
diff --git a/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/dygraph_sharding_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/dygraph_sharding_optimizer.py
@@ -535,6 +535,7 @@ def __init__(self, optimizer, hcg):
 
         # Setting pipeline parallelism overlap
         self.pp_overlap = pp_config.sharding_comm_overlap
+        self.pp_release_grads = pp_config.release_gradients
 
         # TODO(liuzhenhai):support it latter
         assert not self.comm_overlap, "not supported yet"
@@ -604,14 +605,16 @@ def _build_comm_buffers(self, acc_steps, group_size=256 * 1024 * 1024):
                 comm_group,
                 acc_steps,
                 act=HOOK_ACTION.REDUCE_SCATTER,
+                release_grads=self.pp_release_grads,
             )
             self._comm_buffer_list.append(buffer)
 
     def clear_grad(self, set_to_zero=True):
         """
         should clear grad for all parameters in model
         """
-        assert set_to_zero, "should not erase grad buffer"
+        if not self.pp_release_grads:
+            assert set_to_zero, "should not erase grad buffer"
 
         def clear_grad_func(p):
             if hasattr(p, "main_grad") and p.main_grad is not None:
@@ -634,6 +637,10 @@ def clear_grad_func(p):
         for p in self._parameter_list:
             clear_grad_func(p)
 
+        if self.pp_release_grads and not self.pp_overlap:
+            for comm_buffer in self._comm_buffer_list:
+                comm_buffer._clear_grad_storage()
+
     def filter_parameters(self, parameter_list, hcg):
         parameter_list = [
             self._slice_params[param.name] for param in parameter_list
@@ -648,6 +655,10 @@ def reduce_gradients(self, parameter_list, hcg):
         logger.debug("sharding start gradients sync")
         with framework.no_grad():
             for comm_buffer in self._comm_buffer_list:
+                if self.pp_release_grads and comm_buffer.grad_storage is None:
+                    for param in comm_buffer.params:
+                        comm_buffer._copy_grad_to_buffer(param)
+
                 if not self.comm_overlap:
                     comm_buffer._comm_grads()
 
@@ -713,6 +724,9 @@ def _assign_slice_grad(self):
             for param in comm_buffer.params:
                 assert param.name in self._slice_params
                 slice_param = self._slice_params[param.name]
+                if self.pp_release_grads and hasattr(slice_param, "main_grad"):
+                    assert not slice_param.main_grad._is_initialized()
+                    del slice_param.main_grad
                 comm_buffer.assign_slice_grad(param, slice_param)
 
         assert param_num == len(self._parameter_list)
diff --git a/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py b/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py
@@ -208,6 +208,9 @@ def __init__(self, layers, hcg, strategy):
         self._enable_timer = self._strategy.hybrid_configs[
             "pp_configs"
         ].enable_timer
+        self._release_gradients = self._strategy.hybrid_configs[
+            "pp_configs"
+        ].release_gradients
 
         self._sharding_split_param = self._strategy.hybrid_configs[
             "sharding_configs"
@@ -372,7 +375,13 @@ def fused_gradient(
 
                 for group_idx, parameters in var_groups.items():
                     buffer = FusedCommBuffer(
-                        group_idx, parameters, comm_group, acc_steps, act, dst
+                        group_idx,
+                        parameters,
+                        comm_group,
+                        acc_steps,
+                        act,
+                        dst,
+                        release_grads=self._release_gradients,
                     )
                     self._chunk_2_comm_buffers[chunk_idx].append(buffer)
 
@@ -862,7 +871,14 @@ def _optimizer_step(self):
         else:
             self.optimizer.step()
 
-        self.optimizer.clear_grad()
+        if self._release_gradients:
+            self.optimizer.clear_grad(set_to_zero=False)
+            for _, buffers in self._chunk_2_comm_buffers.items():
+                for buffer in buffers:
+                    buffer._clear_grad_storage()
+        else:
+            self.optimizer.clear_grad()
+
         if self.lr_scheduler:
             self.lr_scheduler.step()
 
diff --git a/python/paddle/distributed/fleet/utils/tensor_fusion_helper.py b/python/paddle/distributed/fleet/utils/tensor_fusion_helper.py

Original file line number	Diff line number	Diff line change
`@@ -81,6 +81,7 @@ message PpConfig {`
`81`	`81`	`optional bool enable_timer = 3 [ default = false ];`
`82`	`82`	`optional bool sharding_comm_overlap = 4 [ default = false ];`
`83`	`83`	`optional bool profiling = 5 [ default = false ];`
	`84`	`+ optional bool release_gradients = 6 [ default = false ];`
`84`	`85`	`}`
`85`	`86`
`86`	`87`	`message DygraphShardingConfig {`