Enable recompute allgather in sequence parallel to reduce the memory usage. (#64244)

Xreki · web-flow · commit e8cedebea72a · 2024-05-24T14:14:48.000+08:00
* Optimize the memory usage for sequence parallel.

* Add an option to control whether to recompute allgather in sp.

* Add unittest.
diff --git a/paddle/fluid/framework/distributed_strategy.proto b/paddle/fluid/framework/distributed_strategy.proto
@@ -61,18 +61,20 @@ message ShardingConfig {
 
 // for dygraph
 message MpConfig {
-    optional bool sync_param= 1 [ default = true ];
-    optional bool sync_grad= 2 [ default = false ];
-    optional bool sync_moment= 3 [ default = false ];
-    optional string sync_mode= 4 [ default = 'broadcast' ];
+    optional bool sync_param = 1 [ default = true ];
+    optional bool sync_grad = 2 [ default = false ];
+    optional bool sync_moment = 3 [ default = false ];
+    optional string sync_mode = 4 [ default = 'broadcast' ];
     // Support all_reduce(dx) overlap with matmul(dw) in ColumnParallelLinear
-    optional bool mp_async_allreduce= 5 [default = false ];
+    optional bool mp_async_allreduce = 5 [default = false ];
     // Support skip c_identity in ColumnParallelLinear and RowParallelLinear. Only works when mp_async_allreduce is true.
-    optional bool mp_skip_c_identity= 6 [default = false ];
+    optional bool mp_skip_c_identity = 6 [default = false ];
     // Support fused_linear_param_grad_add in ColumnParallelLinear. Only works when mp_async_allreduce is true.
-    optional bool mp_fused_linear_param_grad_add= 7 [default = false ];
+    optional bool mp_fused_linear_param_grad_add = 7 [default = false ];
     // Broadcast mp input data
-    optional bool need_broadcast_data=8 [default = true];
+    optional bool need_broadcast_data = 8 [default = true];
+    // Recompute AllGather in forward of ColumnSequenceParallelLinear to reduce the memory usage.
+    optional bool recompute_allgather = 9 [default = false];
 }
 
 message PpConfig {
diff --git a/python/paddle/distributed/fleet/utils/sequence_parallel_utils.py b/python/paddle/distributed/fleet/utils/sequence_parallel_utils.py
@@ -237,6 +237,21 @@ def is_fused_linear_param_grad_add_supported():
 _raise_cuda_env_unset_warning_for_sp = True
 
 
+def _check_environment_for_overlap():
+    if int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0")) != 1:
+        global _raise_cuda_env_unset_warning_for_sp
+        if _raise_cuda_env_unset_warning_for_sp:
+            logger.warning(
+                "You set mp_async_allreduce=True or recompute_allgather=True, but you forget to set environment "
+                "variable CUDA_DEVICE_MAX_CONNECTIONS=1, which may leads to performance "
+                "loss. Try to export CUDA_DEVICE_MAX_CONNECTIONS=1 for better performance."
+            )
+        _raise_cuda_env_unset_warning_for_sp = False
+
+        # Using small operation to preempt GPU SMs for all_gather or reduce_scatter to achieve overlap.
+        tmp = paddle.ones([512])
+
+
 class SPInnerOverlapLinear(paddle.autograd.PyLayer):
     @staticmethod
     def forward(
@@ -245,16 +260,22 @@ def forward(
         weight,
         bias,
         fuse_matmul_bias,
+        recompute_allgather,
         mp_fused_linear_param_grad_add,
         model_parallel_group,
     ):
+        ctx.recompute_allgather = recompute_allgather
         ctx.mp_fused_linear_param_grad_add = mp_fused_linear_param_grad_add
         ctx.model_parallel_group = model_parallel_group
 
         world_size = model_parallel_group.nranks
         input_parallel = all_gather(x)
 
-        ctx.save_for_backward(x, weight, bias, input_parallel)
+        if not recompute_allgather:
+            ctx.save_for_backward(x, weight, bias, input_parallel)
+        else:
+            ctx.save_for_backward(x, weight, bias)
+
         if not fuse_matmul_bias:
             output = paddle._C_ops.linear(input_parallel, weight, bias)
         else:
@@ -265,9 +286,26 @@ def forward(
 
     @staticmethod
     def backward(ctx, dy):
-        x, weight, bias, input_parallel = ctx.saved_tensor()
-        parallelism = ctx.model_parallel_group.nranks
+        group = ctx.model_parallel_group
+        parallelism = group.nranks
+
+        if not ctx.recompute_allgather:
+            x, weight, bias, input_parallel = ctx.saved_tensor()
+        else:
+            x, weight, bias = ctx.saved_tensor()
+
+            # all-gather x
+            input_parallel_shape = x.shape
+            input_parallel_shape[0] = input_parallel_shape[0] * parallelism
+            input_parallel = paddle.empty(
+                shape=input_parallel_shape, dtype=x.dtype
+            )
+            allgather_task = dist.all_gather(
+                input_parallel, x, group=group, sync_op=False
+            )
 
+        # compute dx
+        _check_environment_for_overlap()
         if dy.dtype == weight.dtype:
             dinput_parallel = paddle.matmul(dy, weight, transpose_y=True)
         else:
@@ -279,30 +317,24 @@ def backward(ctx, dy):
             dinput_parallel.shape[0] % parallelism == 0
         ), f"Input sequence length {dinput_parallel.shape[0]} can't be divided exactly by sequence parallelism {parallelism}"
 
+        if ctx.recompute_allgather:
+            # wait the finish of all-gather of x
+            allgather_task.wait()
+
+        # reduce-scatter dx
         dx_shape = dinput_parallel.shape
         dx_shape[0] = dx_shape[0] // parallelism
         dx = paddle.empty(shape=dx_shape, dtype=dinput_parallel.dtype)
-        hcg = fleet.get_hybrid_communicate_group()
-        group = hcg.get_model_parallel_group()
         task = dist.stream.reduce_scatter(
             dx,
             dinput_parallel,
             op=dist.ReduceOp.SUM,
             group=group,
             sync_op=False,
         )
-        # Using small operation to preempt GPU SMs for all_reduce to achieve overlap.
-        if int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0")) != 1:
-            global _raise_cuda_env_unset_warning_for_sp
-            if _raise_cuda_env_unset_warning_for_sp:
-                logger.warning(
-                    "You set mp_async_allreduce=True, but you forget to set environment "
-                    "variable CUDA_DEVICE_MAX_CONNECTIONS=1, which may leads to performance "
-                    "loss. Try to export CUDA_DEVICE_MAX_CONNECTIONS=1 for better performance."
-                )
-            _raise_cuda_env_unset_warning_for_sp = False
-            tmp = paddle.ones([512])
 
+        # compute dw and dbias
+        _check_environment_for_overlap()
         if ctx.mp_fused_linear_param_grad_add:
             if not is_fused_linear_param_grad_add_supported():
                 raise NotImplementedError(
@@ -485,21 +517,23 @@ def __init__(
             "mp_configs"
         ]
         self.mp_async_allreduce = mp_configs.mp_async_allreduce
+        self.recompute_allgather = mp_configs.recompute_allgather
 
         self.mp_fused_linear_param_grad_add = (
             self.mp_async_allreduce
             and mp_configs.mp_fused_linear_param_grad_add
         )
 
     def forward(self, x):
-        # sequence parallelism is same as model parallelis, if sequence parallel is true, input shape is [s, b, h],else input shape is [b, s, h]
-        # reuse mp_async_allreduce to do sequence parallelism overlap
+        # sequence parallel is same as tensor parallel, if sequence parallel is true, input shape is [s, b, h], else input shape is [b, s, h]
+        # reuse mp_async_allreduce to do sequence parallel overlap
         if self.mp_async_allreduce:
             output = SPInnerOverlapLinear.apply(
                 x,
                 self.weight,
                 self.bias,
                 self.fuse_matmul_bias,
+                self.recompute_allgather,
                 self.mp_fused_linear_param_grad_add,
                 self.model_parallel_group,
             )
diff --git a/test/collective/fleet/hybrid_parallel_mp_model_with_sequence_parallel.py b/test/collective/fleet/hybrid_parallel_mp_model_with_sequence_parallel.py
@@ -396,7 +396,7 @@ def setUp(self):
         fleet.init(is_collective=True, strategy=strategy)
 
 
-class TestDistSPTraining(unittest.TestCase):
+class TestDistSPTrainingBase(unittest.TestCase):
     def setUp(self):
         strategy = fleet.DistributedStrategy()
         self.model_parallel_size = 2
@@ -408,7 +408,7 @@ def setUp(self):
         }
         fleet.init(is_collective=True, strategy=strategy)
 
-    def train_batch(self, batch, model, optimizer, is_mp):
+    def train_batch(self, batch, model, optimizer):
         output = model(batch)
         loss = output.mean()
         loss.backward()  # do backward
@@ -433,7 +433,7 @@ def build_model_optimizer(self):
         np_fc1 = np.random.random_sample((hidden_size, inner_size))
         np_fc2 = np.random.random_sample((inner_size, hidden_size))
 
-        model_a = SimpleSPNet(
+        sp_model = SimpleSPNet(
             vocab_size,
             hidden_size,
             inner_size,
@@ -442,23 +442,23 @@ def build_model_optimizer(self):
             np_fc2,
             mp_id,
         )
-        optimizer_a = self.build_optimizer(model_a)
-        model_a = fleet.distributed_model(model_a)
-        optimizer_a = fleet.distributed_optimizer(optimizer_a)
+        sp_optimizer = self.build_optimizer(sp_model)
+        sp_model = fleet.distributed_model(sp_model)
+        sp_optimizer = fleet.distributed_optimizer(sp_optimizer)
 
-        model_b = SimpleDPNet(
+        dp_model = SimpleDPNet(
             vocab_size, hidden_size, inner_size, output_size, np_fc1, np_fc2
         )
-        optimizer_b = self.build_optimizer(model_b)
+        dp_optimizer = self.build_optimizer(dp_model)
 
-        return model_a, optimizer_a, model_b, optimizer_b
+        return sp_model, sp_optimizer, dp_model, dp_optimizer
 
-    def test_mp_model(self):
+    def test_sp_accuracy(self):
         (
-            model_a,
-            optimizer_a,
-            model_b,
-            optimizer_b,
+            sp_model,
+            sp_optimizer,
+            dp_model,
+            dp_optimizer,
         ) = self.build_model_optimizer()
 
         for _ in range(5):
@@ -471,15 +471,15 @@ def test_mp_model(self):
                 ),
             )
             batch = paddle.to_tensor(np_data)
-            loss_a = self.train_batch(batch, model_a, optimizer_a, True)
-            loss_b = self.train_batch(batch, model_b, optimizer_b, False)
+            sp_loss = self.train_batch(batch, sp_model, sp_optimizer)
+            dp_loss = self.train_batch(batch, dp_model, dp_optimizer)
 
             np.testing.assert_allclose(
-                loss_a.numpy(), loss_b.numpy(), rtol=1e-5, atol=1e-5
+                sp_loss.numpy(), dp_loss.numpy(), rtol=1e-5, atol=1e-5
             )
 
 
-class TestDistSPTraining2(TestDistSPTraining):
+class TestDistSPTrainingWithConfigs(TestDistSPTrainingBase):
     def setUp(self):
         strategy = fleet.DistributedStrategy()
         self.model_parallel_size = 2
@@ -495,42 +495,8 @@ def setUp(self):
         }
         fleet.init(is_collective=True, strategy=strategy)
 
-    def build_model_optimizer(self):
-        hcg = fleet.get_hybrid_communicate_group()
-        word_size = hcg.get_model_parallel_world_size()
-        mp_id = hcg.get_model_parallel_rank()
-        dp_id = hcg.get_data_parallel_rank()
-        rank_id = dist.get_rank()
-        set_random_seed(1024, dp_id, rank_id)
-
-        np_fc1 = np.random.random_sample((hidden_size, inner_size))
-        np_fc2 = np.random.random_sample((inner_size, hidden_size))
 
-        model_a = SimpleSPNet(
-            vocab_size,
-            hidden_size,
-            inner_size,
-            output_size,
-            np_fc1,
-            np_fc2,
-            mp_id,
-        )
-        model_a = MixPrecisionLayer(model_a)
-        optimizer_a = self.build_optimizer(model_a)
-        optimizer_a = MixPrecisionOptimizer(optimizer_a)
-
-        model_a = fleet.distributed_model(model_a)
-        optimizer_a = fleet.distributed_optimizer(optimizer_a)
-
-        model_b = SimpleDPNet(
-            vocab_size, hidden_size, inner_size, output_size, np_fc1, np_fc2
-        )
-        optimizer_b = self.build_optimizer(model_b)
-
-        return model_a, optimizer_a, model_b, optimizer_b
-
-
-class TestDistSPTraining3(TestDistSPTraining):
+class TestDistSPTrainingAmpWithConfigs(TestDistSPTrainingBase):
     def setUp(self):
         strategy = fleet.DistributedStrategy()
         self.model_parallel_size = 2
@@ -542,6 +508,7 @@ def setUp(self):
             "mp_configs": {
                 "mp_async_allreduce": True,
                 "mp_fused_linear_param_grad_add": True,
+                "recompute_allgather": True,
             },
         }
         fleet.init(is_collective=True, strategy=strategy)
@@ -557,7 +524,7 @@ def build_model_optimizer(self):
         np_fc1 = np.random.random_sample((hidden_size, inner_size))
         np_fc2 = np.random.random_sample((inner_size, hidden_size))
 
-        model_a = SimpleSPNet(
+        sp_model = SimpleSPNet(
             vocab_size,
             hidden_size,
             inner_size,
@@ -566,17 +533,19 @@ def build_model_optimizer(self):
             np_fc2,
             mp_id,
         )
-        optimizer_a = self.build_optimizer(model_a)
+        sp_model = MixPrecisionLayer(sp_model)
+        sp_optimizer = self.build_optimizer(sp_model)
+        sp_optimizer = MixPrecisionOptimizer(sp_optimizer)
 
-        model_a = fleet.distributed_model(model_a)
-        optimizer_a = fleet.distributed_optimizer(optimizer_a)
+        sp_model = fleet.distributed_model(sp_model)
+        sp_optimizer = fleet.distributed_optimizer(sp_optimizer)
 
-        model_b = SimpleDPNet(
+        dp_model = SimpleDPNet(
             vocab_size, hidden_size, inner_size, output_size, np_fc1, np_fc2
         )
-        optimizer_b = self.build_optimizer(model_b)
+        dp_optimizer = self.build_optimizer(dp_model)
 
-        return model_a, optimizer_a, model_b, optimizer_b
+        return sp_model, sp_optimizer, dp_model, dp_optimizer
 
 
 class SimpleSPNetWithoutBias(paddle.nn.Layer):
