[PIR save load] Deal with bf16 in save load function (#63539)

changeyoung98 · web-flow · commit bafb836357b8 · 2024-04-16T16:51:34.000+08:00
* first commit

* refine

* refine merge

* fix bug

* remove decorater

* add ir_gurad

* resolve review comments

* deal with bf16

* refine

* fix

* fix py3
diff --git a/paddle/fluid/pir/serialize_deserialize/src/save_load_parameters.cc b/paddle/fluid/pir/serialize_deserialize/src/save_load_parameters.cc
@@ -18,9 +18,35 @@ limitations under the License. */
 #include "glog/logging.h"
 #include "paddle/fluid/framework/lod_tensor.h"
 #include "paddle/phi/common/port.h"
+#include "paddle/phi/kernels/funcs/data_type_transform.h"
 
 namespace pir {
 
+const phi::DeviceContext* GetDeviceContext(const phi::DenseTensor& x) {
+  phi::DeviceContextPool& pool = phi::DeviceContextPool::Instance();
+  const phi::DeviceContext* dev_ctx = nullptr;
+  auto place = x.place();
+  dev_ctx = pool.Get(place);
+  return dev_ctx;
+}
+
+const phi::DenseTensor CastTensorType(const phi::DeviceContext* dev_ctx,
+                                      const phi::DenseTensor& x,
+                                      phi::DataType out_dtype) {
+  auto place = x.place();
+  if (paddle::platform::is_cpu_place(place)) {
+    auto out = phi::funcs::TransDataType(
+        reinterpret_cast<const phi::CPUContext&>(*dev_ctx), x, out_dtype);
+    return out;
+  } else if (paddle::platform::is_gpu_place(place)) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+    return phi::funcs::TransDataType(
+        reinterpret_cast<const phi::GPUContext&>(*dev_ctx), x, out_dtype);
+#endif
+  }
+  return x;
+}
+
 void SaveFunction(const phi::DenseTensor& x,
                   const std::string& name,
                   const std::string& file_path,
@@ -42,8 +68,16 @@ void SaveFunction(const phi::DenseTensor& x,
       true,
       phi::errors::Unavailable("Cannot open %s to save variables.", file_path));
 
-  paddle::framework::SerializeToStream(fout, x);
-  // TODO(changeyoung98): fp16
+  auto in_dtype = x.dtype();
+  auto out_dtype = save_as_fp16 ? phi::DataType::FLOAT16 : in_dtype;
+
+  const phi::DeviceContext* dev_ctx = GetDeviceContext(x);
+  if (in_dtype != out_dtype) {
+    auto out = CastTensorType(dev_ctx, x, out_dtype);
+    paddle::framework::SerializeToStream(fout, out, *dev_ctx);
+  } else {
+    paddle::framework::SerializeToStream(fout, x, *dev_ctx);
+  }
   fout.close();
   VLOG(6) << "save func done ";
 }
@@ -65,15 +99,28 @@ void SaveCombineFunction(const std::vector<const phi::DenseTensor*>& x,
   MkDirRecursively(DirName(file_path).c_str());
   VLOG(6) << "save func save path: " << file_path;
   std::ostringstream ss;
+  PADDLE_ENFORCE_GT(x.size(),
+                    0UL,
+                    phi::errors::InvalidArgument(
+                        "The number of variables to be saved is %d, expect "
+                        "it to be greater than 0.",
+                        x.size()));
+  const phi::DeviceContext* dev_ctx = GetDeviceContext(*(x[0]));
   for (size_t i = 0; i < x.size(); i++) {
     auto& tensor = *(x[i]);
     PADDLE_ENFORCE_EQ(
         tensor.IsInitialized(),
         true,
         phi::errors::InvalidArgument(
             "The Tensor with Index (%d) to be saved is not initialized.", i));
-    // TODO(changeyoung98): fp16
-    paddle::framework::SerializeToStream(ss, tensor);
+    auto in_dtype = tensor.dtype();
+    auto out_dtype = save_as_fp16 ? phi::DataType::FLOAT16 : in_dtype;
+    if (in_dtype != out_dtype) {
+      auto out = CastTensorType(dev_ctx, tensor, out_dtype);
+      paddle::framework::SerializeToStream(ss, out, *dev_ctx);
+    } else {
+      paddle::framework::SerializeToStream(ss, tensor, *dev_ctx);
+    }
   }
   MkDirRecursively(DirName(file_path).c_str());
   std::ofstream fout(file_path, std::ios::binary);
@@ -101,22 +148,24 @@ void LoadFunction(const std::string& file_path,
   PADDLE_ENFORCE_NOT_NULL(out,
                           phi::errors::InvalidArgument(
                               "The variable to be loaded cannot be found."));
+  const phi::DeviceContext* dev_ctx = GetDeviceContext(*out);
 
   if (seek != -1) {
     PADDLE_ENFORCE_GE(seek,
                       0,
                       phi::errors::InvalidArgument(
                           "seek with tensor must great than or equal to 0"));
-    paddle::platform::DeviceContextPool& pool =
-        paddle::platform::DeviceContextPool::Instance();
-    const paddle::platform::DeviceContext* dev_ctx = nullptr;
-    dev_ctx = pool.Get(paddle::platform::CPUPlace());
     paddle::framework::DeserializeFromStream(fin, out, *dev_ctx, seek, shape);
   } else {
     paddle::framework::DeserializeFromStream(fin, out);
   }
 
-  // TODO(changeyoung98): fp16
+  auto in_dtype = out->dtype();
+  auto out_dtype = load_as_fp16 ? phi::DataType::FLOAT16 : in_dtype;
+  if (in_dtype != out_dtype) {
+    auto cast_in = *out;
+    *out = CastTensorType(dev_ctx, cast_in, out_dtype);
+  }
 }
 
 void LoadCombineFunction(const std::string& file_path,
@@ -130,17 +179,31 @@ void LoadCombineFunction(const std::string& file_path,
                         "Load operator fail to open file %s, please check "
                         "whether the model file is complete or damaged.",
                         file_path));
+
+  PADDLE_ENFORCE_GT(out->size(),
+                    0UL,
+                    phi::errors::InvalidArgument(
+                        "The number of variables to be saved is %d, expect "
+                        "it to be greater than 0.",
+                        out->size()));
+  const phi::DeviceContext* dev_ctx = GetDeviceContext(*(out->at(0)));
   for (size_t i = 0; i < names.size(); i++) {
     auto tensor = out->at(i);
     paddle::framework::DeserializeFromStream(fin, tensor);
+
+    auto in_dtype = tensor->dtype();
+    auto out_dtype = load_as_fp16 ? phi::DataType::FLOAT16 : in_dtype;
+    if (in_dtype != out_dtype) {
+      auto cast_in = *tensor;
+      *tensor = CastTensorType(dev_ctx, cast_in, out_dtype);
+    }
   }
   fin.peek();
   PADDLE_ENFORCE_EQ(
       fin.eof(),
       true,
       phi::errors::Unavailable("Not allowed to load partial data via "
                                "load_combine_op, please use load_op instead."));
-  // TODO(changeyoung98): fp16
 }
 
 }  // namespace pir
diff --git a/test/ir/pir/test_save_load_params.py b/test/ir/pir/test_save_load_params.py
@@ -19,7 +19,6 @@
 import numpy as np
 
 import paddle
-from paddle import base
 from paddle.optimizer import Adam
 from paddle.pir_utils import IrGuard
 
@@ -61,33 +60,21 @@ def test_params1(self):
                 main_program, paddle.static.Program()
             ):
                 x = paddle.static.data(
-                    'x', shape=[-1, 16], dtype='float32', lod_level=1
-                )
-
-                linear = paddle.nn.Linear(
-                    in_features=x.shape[-1],
-                    out_features=16,
-                    weight_attr=paddle.nn.initializer.UniformInitializer(
-                        low=-0.5,
-                        high=0.5,
-                        seed=10,
-                        diag_num=16,
-                        diag_step=16,
-                        diag_val=1.0,
-                    ),
-                )
-                y = linear(x)
-
-                place = base.CPUPlace()
-                x_tensor = base.create_lod_tensor(
-                    np.random.rand(3, 16).astype("float32"), [[1, 2]], place
+                    name="static_x", shape=[None, IMAGE_SIZE], dtype='float32'
                 )
-                exe = base.Executor(place)
+                z = paddle.static.nn.fc(x, 10)
+                z = paddle.static.nn.fc(z, 10, bias_attr=False)
+                loss = paddle.mean(z)
+                opt = Adam(learning_rate=1e-3)
+                opt.minimize(loss)
+                place = paddle.CPUPlace()
+                exe = paddle.static.Executor(place)
                 exe.run(paddle.static.default_startup_program())
-                ret = exe.run(
+                fake_inputs = np.random.randn(2, IMAGE_SIZE).astype('float32')
+                exe.run(
                     main_program,
-                    feed={'x': x_tensor},
-                    fetch_list=[y],
+                    feed={'static_x': fake_inputs},
+                    fetch_list=[loss],
                 )
                 scope = paddle.static.global_scope()
                 params = main_program.global_block().all_parameters()
@@ -137,10 +124,11 @@ def test_params2(self):
                 save_dir = os.path.join(self.temp_dir.name, "save_params")
                 for k, v in param_dict.items():
                     path = os.path.join(save_dir, k, '.pdparams')
-                    paddle.base.core.save_func(v, k, path, True, False)
+                    # test fp16
+                    paddle.base.core.save_func(v, k, path, True, True)
                     tensor = param_dict[k]
                     tensor.set(np.zeros_like(np.array(tensor)), place)
-                    paddle.base.core.load_func(path, -1, [], False, tensor)
+                    paddle.base.core.load_func(path, -1, [], True, tensor)
                     np.testing.assert_array_equal(tensor, v)
 
                 for k, v in opt_dict.items():
@@ -172,6 +160,14 @@ def test_params2(self):
                 paddle.base.core.save_combine_func(
                     param_vec, list(param_dict.keys()), path, True, False, True
                 )
+                # save as fp16
+                paddle.base.core.save_combine_func(
+                    param_vec, list(param_dict.keys()), path, True, True, False
+                )
+                # load as fp16
+                paddle.base.core.load_combine_func(
+                    path, list(param_dict.keys()), param_new, True
+                )
 
                 # test save_vars
                 path_prefix = os.path.join(save_dir, 'new')
