[AutoParallel] Support FusedRoPE shard on seq_len in Semi-Auto (#62053)

MarioLulab · web-flow · commit 928356f77344 · 2024-03-05T17:23:45.000+08:00
* add rotary_emb_base attr * add testcases * fix cc test * remove annotation * modify spmd * polish code * Revert "add testcases" This reverts commit 6caec37. * Revert "add rotary_emb_base attr" This reverts commit 441b816. * Revert "fix cc test" This reverts commit 68457d1. * add placements checking
diff --git a/paddle/phi/infermeta/spmd_rules/fused_rope.cc b/paddle/phi/infermeta/spmd_rules/fused_rope.cc
@@ -80,8 +80,9 @@ void check_k_or_v(const DistMetaTensor& k_or_v,
 void check_sin_cos(const DistMetaTensor& sin,
                    const DistMetaTensor& cos,
                    const DistMetaTensor& position_ids,
-                   const std::vector<int64_t>& q_shape,
-                   bool time_major) {
+                   const int64_t batch_size,
+                   const int64_t seq_len,
+                   const int64_t head_dim) {
   PADDLE_ENFORCE_EQ(sin.dims(),
                     cos.dims(),
                     phi::errors::InvalidArgument(
@@ -98,13 +99,6 @@ void check_sin_cos(const DistMetaTensor& sin,
       phi::errors::InvalidArgument(
           "The Tensor sin/cos's ndim must be 2 or 4. but given [%d]", ndim));
 
-  const int kBatchDimIndex = time_major ? 1 : 0;
-  const int kSeqlenDimIndex = time_major ? 0 : 1;
-
-  int batch_size = q_shape[kBatchDimIndex];
-  int seq_len = q_shape[kSeqlenDimIndex];
-  int head_dim = q_shape[kHeadDimIndex];
-
   int seq_len_dim_index = ndim == 2 ? 0 : 1;
   int head_dim_index = ndim == 2 ? 1 : 3;
   if (ndim == 4) {
@@ -143,9 +137,10 @@ void check_sin_cos(const DistMetaTensor& sin,
         phi::errors::InvalidArgument(
             "The batch_size and seq_len of position_ids must be the same as "
             "those of q. But received position_ids's "
-            "shape is {%s}, q's shape is {%s}.",
+            "shape is {%s}, q's batch_size is {%d}, q's seq_len is {%d}.",
             str_join(position_ids_shape),
-            str_join(q_shape)));
+            batch_size,
+            seq_len));
   } else {
     PADDLE_ENFORCE_EQ(
         (shape[seq_len_dim_index] == seq_len &&
@@ -162,8 +157,10 @@ void check_sin_cos(const DistMetaTensor& sin,
 void infer_sin_cos(const DistMetaTensor& sin,
                    const DistMetaTensor& cos,
                    const DistMetaTensor& position_ids,
+                   const TensorDistAttr& q_dist_attr_dst,
                    const std::vector<int64_t>& q_shape,
                    bool time_major,
+                   bool enable_sequence_parallel,
                    TensorDistAttr* sin_dist_attr_dst,
                    TensorDistAttr* cos_dist_attr_dst) {
   const TensorDistAttr& sin_dist_attr_src = sin.dist_attr();
@@ -178,13 +175,39 @@ void infer_sin_cos(const DistMetaTensor& sin,
   // if one of sin cos is empty, they are all useless in kernel
   if (!IsEmpty(sin_shape) && !IsEmpty(cos_shape)) {
     // check sin, cos, position_ids's shape
-    check_sin_cos(sin, cos, position_ids, q_shape, time_major);
-    if (sin_shape.size() == 4) {
-      *sin_dist_attr_dst = UnShardTensorDims(sin_dist_attr_src, {1, 3});
-      *cos_dist_attr_dst = UnShardTensorDims(cos_dist_attr_src, {1, 3});
-    } else {
-      *sin_dist_attr_dst = UnShardTensorDims(sin_dist_attr_src, {0, 1});
-      *cos_dist_attr_dst = UnShardTensorDims(cos_dist_attr_src, {0, 1});
+    const int kBatchDimIndex = time_major ? 1 : 0;
+    const int kSeqlenDimIndex = time_major ? 0 : 1;
+    int batch_size = q_shape[kBatchDimIndex];
+    int seq_len = q_shape[kSeqlenDimIndex];
+    int head_dim = q_shape[kHeadDimIndex];
+
+    int seq_len_dim_index = sin_shape.size() == 4 ? 1 : 0;
+    int head_dim_index = sin_shape.size() == 4 ? 3 : 1;
+
+    check_sin_cos(sin, cos, position_ids, batch_size, seq_len, head_dim);
+
+    *sin_dist_attr_dst =
+        enable_sequence_parallel
+            ? UnShardTensorDims(sin_dist_attr_src, {head_dim_index})
+            : UnShardTensorDims(sin_dist_attr_src,
+                                {seq_len_dim_index, head_dim_index});
+    *cos_dist_attr_dst =
+        enable_sequence_parallel
+            ? UnShardTensorDims(sin_dist_attr_src, {head_dim_index})
+            : UnShardTensorDims(cos_dist_attr_src,
+                                {seq_len_dim_index, head_dim_index});
+
+    if (enable_sequence_parallel) {
+      // shard on seq_len dimension
+      std::vector<int64_t> sin_dims_mapping = sin_dist_attr_dst->dims_mapping();
+      sin_dims_mapping[seq_len_dim_index] =
+          q_dist_attr_dst.dims_mapping()[kSeqlenDimIndex];
+      sin_dist_attr_dst->set_dims_mapping(sin_dims_mapping);
+
+      std::vector<int64_t> cos_dims_mapping = cos_dist_attr_dst->dims_mapping();
+      cos_dims_mapping[seq_len_dim_index] =
+          q_dist_attr_dst.dims_mapping()[kSeqlenDimIndex];
+      cos_dist_attr_dst->set_dims_mapping(cos_dims_mapping);
     }
   }
 }
@@ -237,9 +260,24 @@ SpmdInfo FusedRopeInferSpmd(const DistMetaTensor& q,
       GetDimsMappingForAxes(qkv_axes, axis_to_dim_map);
   TensorDistAttr q_dist_attr_dst = CopyTensorDistAttrForOutput(q_dist_attr_src);
   q_dist_attr_dst.set_dims_mapping(out_dims_mapping);
+
   const int kSeqlenDimIndex = time_major ? 0 : 1;
-  q_dist_attr_dst =
-      UnShardTensorDims(q_dist_attr_dst, {kSeqlenDimIndex, kHeadDimIndex});
+  // if one of sin cos is empty, they are all useless in kernel
+  bool is_sin_cos_none = IsEmpty(common::vectorize(sin.dims())) ||
+                         IsEmpty(common::vectorize(cos.dims()));
+
+  // Enable sharding on seq_len dimension only if sin/cos is not None and
+  // position_ids is None
+  bool enable_sequence_parallel =
+      !is_sin_cos_none && is_ids_none &&
+      IsDimSharded(q_dist_attr_dst, kSeqlenDimIndex);
+  if (enable_sequence_parallel) {
+    // Sharded along seq_len dimension
+    q_dist_attr_dst = UnShardTensorDims(q_dist_attr_dst, {kHeadDimIndex});
+  } else {
+    q_dist_attr_dst =
+        UnShardTensorDims(q_dist_attr_dst, {kSeqlenDimIndex, kHeadDimIndex});
+  }
 
   TensorDistAttr k_dist_attr_dst = CopyTensorDistAttrForOutput(k_dist_attr_src);
   k_dist_attr_dst.set_process_mesh(q_dist_attr_dst.process_mesh());
@@ -258,8 +296,10 @@ SpmdInfo FusedRopeInferSpmd(const DistMetaTensor& q,
   infer_sin_cos(sin,
                 cos,
                 position_ids,
+                q_dist_attr_dst,
                 q_shape,
                 time_major,
+                enable_sequence_parallel,
                 &sin_dist_attr_dst,
                 &cos_dist_attr_dst);
 
@@ -331,8 +371,24 @@ SpmdInfo FusedRopeInferSpmdReverse(const DistMetaTensor& q,
   q_dist_attr_dst.set_dims_mapping(dims_mapping);
 
   const int kSeqlenDimIndex = time_major ? 0 : 1;
-  q_dist_attr_dst =
-      UnShardTensorDims(q_dist_attr_dst, {kSeqlenDimIndex, kHeadDimIndex});
+  // if one of sin cos is empty, they are all useless in kernel
+  bool is_sin_cos_none = IsEmpty(common::vectorize(sin.dims())) ||
+                         IsEmpty(common::vectorize(cos.dims()));
+  bool is_ids_none = IsEmpty(common::vectorize(position_ids.dims()));
+
+  // Enable sharding on seq_len dimension only if sin/cos is not None and
+  // position_ids is None
+  bool enable_sequence_parallel =
+      !is_sin_cos_none && is_ids_none &&
+      IsDimSharded(q_dist_attr_dst, kSeqlenDimIndex);
+  if (enable_sequence_parallel) {
+    // Sharded along seq_len dimension
+    q_dist_attr_dst = UnShardTensorDims(q_dist_attr_dst, {kHeadDimIndex});
+  } else {
+    q_dist_attr_dst =
+        UnShardTensorDims(q_dist_attr_dst, {kSeqlenDimIndex, kHeadDimIndex});
+  }
+
   TensorDistAttr out_q_dist_attr_dst = q_dist_attr_dst;
 
   TensorDistAttr k_dist_attr_dst = CopyTensorDistAttrForOutput(k.dist_attr());
@@ -356,8 +412,10 @@ SpmdInfo FusedRopeInferSpmdReverse(const DistMetaTensor& q,
   infer_sin_cos(sin,
                 cos,
                 position_ids,
+                out_q_dist_attr_dst,
                 out_q_shape,
                 time_major,
+                enable_sequence_parallel,
                 &sin_dist_attr_dst,
                 &cos_dist_attr_dst);
 
@@ -367,7 +425,6 @@ SpmdInfo FusedRopeInferSpmdReverse(const DistMetaTensor& q,
   TensorDistAttr position_ids_dist_attr_dst =
       CopyTensorDistAttrForOutput(position_ids.dist_attr());
 
-  bool is_ids_none = IsEmpty(common::vectorize(position_ids.dims()));
   if (!is_ids_none) {
     position_ids_dist_attr_dst.set_dims_mapping(position_ids_dims_mapping);
     position_ids_dist_attr_dst =
diff --git a/test/auto_parallel/semi_auto_parallel_for_fused_rope.py b/test/auto_parallel/semi_auto_parallel_for_fused_rope.py
@@ -223,6 +223,65 @@ def test_common_case_time_major(self):
         self.check_tensor_eq(dist_q.grad, q.grad)
         self.check_tensor_eq(dist_k.grad, k.grad)
 
+    def test_common_case_time_major_shard_seq(self):
+        paddle.seed(self._seed)
+        np.random.seed(self._seed)
+        # [seq_len, bs, num_heads, head_dim]
+        qkv_shape = [self._seq_len, self._bs, self._num_heads, self._head_dim]
+        q = paddle.randn(qkv_shape, self._dtype)
+        q.stop_gradient = False
+
+        dist_q = dist.shard_tensor(q, self._mesh, dist.Shard(0))
+        dist_q.stop_gradient = False
+
+        k = paddle.randn(qkv_shape, self._dtype)
+        k.stop_gradient = False
+        dist_k = dist.shard_tensor(k, self._mesh, dist.Shard(2))
+        dist_k.stop_gradient = False
+
+        sin = paddle.randn(self._sin_cos_shape, self._dtype)
+        sin.stop_gradient = True
+        dist_sin = dist.shard_tensor(sin, self._mesh, dist.Replicate())
+        dist_sin.stop_gradient = True
+
+        cos = paddle.randn(self._sin_cos_shape, self._dtype)
+        cos.stop_gradient = True
+        dist_cos = dist.shard_tensor(cos, self._mesh, dist.Replicate())
+        dist_cos.stop_gradient = True
+
+        dist_out_q, dist_out_k, _ = fused_rotary_position_embedding(
+            q=dist_q,
+            k=dist_k,
+            sin=dist_sin,
+            cos=dist_cos,
+            position_ids=None,
+            use_neox_rotary_style=False,
+            time_major=True,
+        )
+        out_q, out_k, _ = fused_rotary_position_embedding(
+            q=q,
+            k=k,
+            sin=sin,
+            cos=cos,
+            position_ids=None,
+            use_neox_rotary_style=False,
+            time_major=True,
+        )
+
+        self.check_placements(dist_out_q, [dist.Shard(0)])
+        self.check_placements(dist_out_k, [dist.Shard(0)])
+
+        self.check_tensor_eq(out_q, dist_out_q)
+        self.check_tensor_eq(out_k, dist_out_k)
+
+        dist_out = dist_out_q + dist_out_k
+        out = out_q + out_k
+        dist_out.backward()
+        out.backward()
+
+        self.check_tensor_eq(dist_q.grad, q.grad)
+        self.check_tensor_eq(dist_k.grad, k.grad)
+
     def run_test_case(self):
         if self._backend == "gpu":
             paddle.set_device("gpu:" + str(dist.get_rank()))
@@ -235,6 +294,7 @@ def run_test_case(self):
         self.test_only_q_input_time_major()
         self.test_common_case()
         self.test_common_case_time_major()
+        self.test_common_case_time_major_shard_seq()
 
 
 if __name__ == '__main__':
