[XPU] fix that TensorCopy changes the shape of dst Tensor in c_softmax_with_cross_entropy_op_xpu (#63780)

dynamicheart · web-flow · commit 172aa4419a06 · 2024-04-24T12:27:56.000+08:00
* [XPU] fix that TensorCopy changes the shape of dst Tensor in c_softmax_with_cross_entropy_op_xpu

* add tests
diff --git a/paddle/fluid/operators/collective/c_softmax_with_cross_entropy_op_xpu.cc b/paddle/fluid/operators/collective/c_softmax_with_cross_entropy_op_xpu.cc
@@ -248,8 +248,11 @@ struct CSoftmaxWithCrossEntropyProcessGroupFunctor<phi::XPUContext, T> {
         N * 1);
     PADDLE_ENFORCE_XDNN_SUCCESS(ret, "sub");
 
-    framework::TensorCopy(
-        softmax_2d, ctx.GetPlace(), ctx.device_context(), softmax);
+    memory::Copy(ctx.GetPlace(),
+                 softmax->data(),
+                 ctx.GetPlace(),
+                 softmax_2d.data(),
+                 N * D * sizeof(T));
   }
 };
 
@@ -510,8 +513,11 @@ struct CSoftmaxWithCrossEntropyFunctor<phi::XPUContext, T> {
         N * 1);
     PADDLE_ENFORCE_XDNN_SUCCESS(ret, "sub");
 
-    framework::TensorCopy(
-        softmax_2d, ctx.GetPlace(), ctx.device_context(), softmax);
+    memory::Copy(ctx.GetPlace(),
+                 softmax->data(),
+                 ctx.GetPlace(),
+                 softmax_2d.data(),
+                 N * D * sizeof(T));
   }
 };
 
diff --git a/test/xpu/collective_softmax_with_cross_entropy_op_xpu.py b/test/xpu/collective_softmax_with_cross_entropy_op_xpu.py
@@ -17,6 +17,8 @@
 import sys
 
 import numpy as np
+
+sys.path.append("../legacy_test")
 from op_test import convert_float_to_uint16
 from test_collective_base_xpu import (
     DataTypeCast,
@@ -34,22 +36,24 @@
 class TestCollectiveSoftmaxWithCE(TestCollectiveRunnerBase):
     def __init__(self):
         self.global_ring_id = 0
-        self.batch_size = 10
+        self.batch_size = 1
+        self.seq_len = 10
         self.num_class = 1000
         self.nranks = 2
         self.ring_id = 0
         self.local_elements = int(self.num_class / self.nranks)
 
+        self.logits_shape = [self.seq_len, self.local_elements]
+        self.label_shape = [self.seq_len, 1]
+
     def get_model(self, main_prog, startup_program, rank):
         with program_guard(main_prog, startup_program):
             logits = data(
                 name="Logits",
-                shape=[self.batch_size, self.local_elements],
+                shape=self.logits_shape,
                 dtype=self.dtype,
             )
-            label = data(
-                name="Label", shape=[self.batch_size, 1], dtype='int32'
-            )
+            label = data(name="Label", shape=self.label_shape, dtype='int32')
             softmax = main_prog.current_block().create_var(
                 name="Softmax",
                 dtype=logits.dtype,
@@ -66,7 +70,7 @@ def get_model(self, main_prog, startup_program, rank):
             )
             loss_grad = main_prog.current_block().create_var(
                 name="Loss@GRAD",
-                shape=[self.batch_size, 1],
+                shape=self.label_shape,
                 dtype=logits.dtype,
                 type=core.VarDesc.VarType.LOD_TENSOR,
                 persistable=False,
@@ -112,6 +116,19 @@ def run_trainer(self, args):
             startup_prog, rank, self.nranks, True, current_endpoint, endpoints
         )
         self.dtype = args["dtype"]
+
+        # if batch_size = 1, we treat logits/labels as 2D tensors
+        # if batch_size > 1, we treat logits/labels as 3D tensors
+        if self.batch_size is not None:
+            self.batch_size = int(args["batch_size"])
+        if self.batch_size > 1:
+            self.logits_shape = [
+                self.batch_size,
+                self.seq_len,
+                self.local_elements,
+            ]
+            self.label_shape = [self.batch_size, self.seq_len, 1]
+
         np_dtype = DataTypeCast(args["dtype"])
         loss, softmax = self.get_model(train_prog, startup_prog, rank)
         device_id = int(os.getenv("FLAGS_selected_xpus", "0"))
@@ -124,12 +141,12 @@ def run_trainer(self, args):
         label = np.random.randint(
             0,
             self.num_class,
-            size=(self.batch_size, 1),
+            size=self.label_shape,
             dtype='int32',
         )
         # use FAKE loss_grad here, only to examine the correctness of grad func
         loss_grad_fp32 = np.random.uniform(
-            low=-10.0, high=10.0, size=(self.batch_size, 1)
+            low=-10.0, high=10.0, size=self.label_shape
         ).astype(np.float32)
         if args["dtype"] == "bfloat16":
             loss_grad = convert_float_to_uint16(loss_grad_fp32)
@@ -139,7 +156,7 @@ def run_trainer(self, args):
         # each xpu uses own half of logits
         np.random.seed(os.getpid())
         logits_fp32 = np.random.uniform(
-            low=-40.0, high=40.0, size=(self.batch_size, self.local_elements)
+            low=-40.0, high=40.0, size=self.logits_shape
         ).astype(np.float32)
         if args["dtype"] == "bfloat16":
             logits = convert_float_to_uint16(logits_fp32)
diff --git a/test/xpu/test_collective_base_xpu.py b/test/xpu/test_collective_base_xpu.py
@@ -167,6 +167,7 @@ def runtime_main(test_class, col_type, sub_type):
     args["currentendpoint"] = os.getenv("PADDLE_CURRENT_ENDPOINT")
     args["col_type"] = col_type
     args["dtype"] = os.getenv("DTYPE")
+    args["batch_size"] = os.getenv("BATCH_SIZE")
     args["dynamic_static_unified_comm"] = bool(
         int(os.getenv("FLAGS_dynamic_static_unified_comm", "0"))
     )
diff --git a/test/xpu/test_collective_softmax_with_cross_entropy_xpu.py b/test/xpu/test_collective_softmax_with_cross_entropy_xpu.py
@@ -59,18 +59,25 @@ def cross_entropy(softmax, label, soft_label, axis, ignore_index=-1):
 
 
 def softmax_with_cross_entropy_grad(softmax, label, loss_grad, axis):
-    logit_grad = softmax.copy()
     shape = softmax.shape
     axis %= len(shape)
     n = int(np.prod(shape[:axis]))
     d = int(np.prod(shape[axis:]))
+    logit_grad_2d = softmax.copy().reshape(n, d)
+    loss_grad_2d = loss_grad.reshape(n, 1)
+    label_2d = label.reshape(n, 1)
     for i in range(n * d):
         row = int(i / d)
         col = i % d
-        if col == label[row]:
-            logit_grad[row][col] = (logit_grad[row][col] - 1.0) * loss_grad[row]
+        if col == label_2d[row]:
+            logit_grad_2d[row][col] = (
+                logit_grad_2d[row][col] - 1.0
+            ) * loss_grad_2d[row]
         else:
-            logit_grad[row][col] = logit_grad[row][col] * loss_grad[row]
+            logit_grad_2d[row][col] = (
+                logit_grad_2d[row][col] * loss_grad_2d[row]
+            )
+    logit_grad = logit_grad_2d.reshape(softmax.shape)
     return logit_grad
 
 
@@ -83,8 +90,9 @@ class TestCSoftmaxWithCEOp(TestDistBase):
         def _setup_config(self):
             pass
 
-        def test_softmax_with_ce(self):
-            self.batch_size = 10
+        def test_softmax_with_ce_2d_logits(self):
+            self.batch_size = 1
+            self.seq_len = 10
             self.num_class = 1000
             self.check_with_place(
                 "collective_softmax_with_cross_entropy_op_xpu.py",
@@ -108,6 +116,7 @@ def check_with_place(
                 "LD_PRELOAD": os.getenv("LD_PRELOAD", ""),
                 "GLOG_v": "3",
                 "DTYPE": dtype,
+                "BATCH_SIZE": str(self.batch_size),
                 "FLAGS_dynamic_static_unified_comm": "0",
             }
             required_envs.update(need_envs)
@@ -120,36 +129,45 @@ def check_with_place(
                 model_file, required_envs
             )
 
+            # if batch_size = 1, we treat logits/labels as 2D tensors
+            # if batch_size > 1, we treat logits/labels as 3D tensors
+            local_elements = int(self.num_class / 2)
+            if self.batch_size > 1:
+                logits_shape = [self.batch_size, self.seq_len, local_elements]
+                label_shape = [self.batch_size, self.seq_len, 1]
+            else:
+                logits_shape = [self.seq_len, local_elements]
+                label_shape = [self.seq_len, 1]
+
             # get data that is shared by both ranks
             np.random.seed(os.getuid())
             label = np.random.randint(
-                0, self.num_class, size=(self.batch_size, 1), dtype='int32'
+                0, self.num_class, size=label_shape, dtype='int32'
             )
             loss_grad = np.random.uniform(
-                low=-10.0, high=10.0, size=(self.batch_size, 1)
+                low=-10.0, high=10.0, size=label_shape
             ).astype(np_dtype)
 
-            local_elements = int(self.num_class / 2)
             # get input data for rank 0
             np.random.seed(pid0)
             input0 = np.random.uniform(
-                low=-40.0, high=40.0, size=(self.batch_size, local_elements)
+                low=-40.0, high=40.0, size=logits_shape
             ).astype(np_dtype)
 
             # get input data for rank 1
             np.random.seed(pid1)
             input1 = np.random.uniform(
-                low=-40.0, high=40.0, size=(self.batch_size, local_elements)
+                low=-40.0, high=40.0, size=logits_shape
             ).astype(np_dtype)
 
             # get combined input data
-            inputs = np.concatenate((input0, input1), axis=1)
+            inputs = np.concatenate((input0, input1), axis=-1)
 
             # calculate analytic result
-            need_softmax = np.apply_along_axis(stable_softmax, 1, inputs)
-            need_loss = cross_entropy(need_softmax, label, False, 1)
+            need_softmax = np.apply_along_axis(stable_softmax, -1, inputs)
+            need_loss = cross_entropy(need_softmax, label, False, -1)
             need_logits_grad = softmax_with_cross_entropy_grad(
-                need_softmax, label, loss_grad, axis=1
+                need_softmax, label, loss_grad, axis=-1
             )
 
             # get real result
@@ -162,8 +180,8 @@ def check_with_place(
                 loss1 = convert_uint16_to_float(loss1)
                 softmax1 = convert_uint16_to_float(softmax1)
                 logits_grad1 = convert_uint16_to_float(logits_grad1)
-            softmax = np.concatenate((softmax0, softmax1), axis=1)
-            logits_grad = np.concatenate((logits_grad0, logits_grad1), axis=1)
+            softmax = np.concatenate((softmax0, softmax1), axis=-1)
+            logits_grad = np.concatenate((logits_grad0, logits_grad1), axis=-1)
 
             # compare results
             rtol = 1e-6
@@ -180,6 +198,20 @@ def check_with_place(
                 logits_grad, need_logits_grad, rtol=rtol, atol=atol
             )
 
+    class TestCSoftmaxWithCEOp1(TestCSoftmaxWithCEOp):
+        def _setup_config(self):
+            pass
+
+        def test_softmax_with_ce_3d_logis(self):
+            self.batch_size = 2
+            self.seq_len = 10
+            self.num_class = 1000
+            self.check_with_place(
+                "collective_softmax_with_cross_entropy_op_xpu.py",
+                "softmax_with_ce",
+                self.in_type_str,
+            )
+
 
 support_types = get_xpu_op_support_types('c_softmax_with_cross_entropy')
 for stype in support_types:

Original file line number	Diff line number	Diff line change
`@@ -167,6 +167,7 @@ def runtime_main(test_class, col_type, sub_type):`
`167`	`167`	`args["currentendpoint"] = os.getenv("PADDLE_CURRENT_ENDPOINT")`
`168`	`168`	`args["col_type"] = col_type`
`169`	`169`	`args["dtype"] = os.getenv("DTYPE")`
	`170`	`+ args["batch_size"] = os.getenv("BATCH_SIZE")`
`170`	`171`	`args["dynamic_static_unified_comm"] = bool(`
`171`	`172`	`int(os.getenv("FLAGS_dynamic_static_unified_comm", "0"))`
`172`	`173`	`)`