diff --git a/paddle/fluid/jit/all.h b/paddle/fluid/jit/all.h
new file mode 100644
index 0000000000000000000000000000000000000000..5a571a72a28243540c35599992170c4b47c9dbbf
--- /dev/null
+++ b/paddle/fluid/jit/all.h
@@ -0,0 +1,20 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "base_function.h"
+#include "layer.h"
+#include "serializer.h"
+#include "serializer_utils.h"
diff --git a/paddle/fluid/jit/base_function.h b/paddle/fluid/jit/base_function.h
index df774d8fd84c701b789d403344945cc8bf1fa950..50dadaf4ae2276b5656811ec9b0fdc12851e75f7 100644
--- a/paddle/fluid/jit/base_function.h
+++ b/paddle/fluid/jit/base_function.h
@@ -15,7 +15,6 @@
 #pragma once
 
 #include "paddle/phi/api/include/tensor.h"
-#include "paddle/phi/core/dense_tensor.h"
 
 namespace paddle {
 namespace jit {
diff --git a/paddle/fluid/jit/compilation_unit.cc b/paddle/fluid/jit/compilation_unit.cc
index 60d42d045b0e3b7cc52587874a4787830d91d4f2..5a434fba176d3769680b1a4ecc31edfa68afec00 100644
--- a/paddle/fluid/jit/compilation_unit.cc
+++ b/paddle/fluid/jit/compilation_unit.cc
@@ -16,6 +16,8 @@
 
 #include "paddle/phi/core/enforce.h"
 
+#include "paddle/fluid/jit/base_function.h"
+
 namespace paddle {
 namespace jit {
 
diff --git a/paddle/fluid/jit/compilation_unit.h b/paddle/fluid/jit/compilation_unit.h
index 45a771b6494015cfc8cb3e7e384b0e5ed6a4a9db..535e92fe88473e85443d0d2a4180393a2dc6623f 100644
--- a/paddle/fluid/jit/compilation_unit.h
+++ b/paddle/fluid/jit/compilation_unit.h
@@ -14,13 +14,14 @@
 
 #pragma once
 
+#include <memory>
 #include <string>
 #include <unordered_map>
-
-#include "paddle/fluid/jit/base_function.h"
+#include <vector>
 
 namespace paddle {
 namespace jit {
+class BaseFunction;
 using Name2FunctionMap =
     std::unordered_map<std::string, std::shared_ptr<BaseFunction>>;
 
diff --git a/paddle/fluid/jit/function_schema.cc b/paddle/fluid/jit/function_schema.cc
index 20cbcfdbd1c88d218578ab54bafc2d733c07a773..8150d3b2e7589c3cfa3543262ec60d74eb410683 100644
--- a/paddle/fluid/jit/function_schema.cc
+++ b/paddle/fluid/jit/function_schema.cc
@@ -14,6 +14,7 @@
 
 #include "paddle/fluid/jit/function_schema.h"
 
+#include "paddle/fluid/framework/program_desc.h"
 #include "paddle/phi/core/enforce.h"
 
 #include "paddle/fluid/jit/function_utils.h"
@@ -52,14 +53,13 @@ void FunctionSchema::AddOutputArg(const std::string& name) {
 FunctionInfo::FunctionInfo(const std::string& func_name,
                            const std::vector<std::string>& param_names,
                            const framework::ProgramDesc& program_desc)
-    : func_name_(func_name),
-      param_names_(param_names),
-      program_desc_(program_desc) {
+    : func_name_(func_name), param_names_(param_names) {
+  program_desc_.reset(new framework::ProgramDesc(program_desc));
   // Parse FunctionSchema
-  for (auto& in_name : program_desc_.GetFeedTargetNames()) {
+  for (auto& in_name : program_desc_->GetFeedTargetNames()) {
     schema_.AddInputArg(in_name);
   }
-  for (auto& out_name : program_desc_.GetFetchTargetNames()) {
+  for (auto& out_name : program_desc_->GetFetchTargetNames()) {
     schema_.AddOutputArg(out_name);
   }
 }
@@ -67,7 +67,7 @@ FunctionInfo::FunctionInfo(const std::string& func_name,
 const std::string& FunctionInfo::FunctionName() const { return func_name_; }
 
 const framework::ProgramDesc& FunctionInfo::ProgramDesc() const {
-  return program_desc_;
+  return *program_desc_.get();
 }
 
 const std::vector<std::string>& FunctionInfo::ParamNames() const {
@@ -83,7 +83,7 @@ const std::vector<std::string> FunctionInfo::OutputArgNames() const {
 }
 
 void FunctionInfo::RemoveDescFeedFetch() {
-  utils::RemoveFeedFetch(&program_desc_);
+  utils::RemoveFeedFetch(program_desc_.get());
 }
 
 }  // namespace jit
diff --git a/paddle/fluid/jit/function_schema.h b/paddle/fluid/jit/function_schema.h
index 5dcea8517e40eea1dfcdfded64548eb113ad6e7f..9f593dd7eee2411f761250a7e8adca9cbc6ba05f 100644
--- a/paddle/fluid/jit/function_schema.h
+++ b/paddle/fluid/jit/function_schema.h
@@ -14,15 +14,17 @@
 
 #pragma once
 
+#include <memory>
 #include <string>
 #include <vector>
 
-#include "paddle/fluid/framework/program_desc.h"
-#include "paddle/fluid/framework/variable.h"
-
 namespace paddle {
+
+namespace framework {
+class ProgramDesc;
+}  // namespace framework
+
 namespace jit {
-using Variable = paddle::framework::Variable;
 
 class Argument {
  public:
@@ -75,7 +77,7 @@ class FunctionInfo {
  private:
   std::string func_name_;
   std::vector<std::string> param_names_;
-  framework::ProgramDesc program_desc_;
+  std::shared_ptr<framework::ProgramDesc> program_desc_;
   FunctionSchema schema_;
 };
 
diff --git a/paddle/fluid/jit/function_utils.cc b/paddle/fluid/jit/function_utils.cc
index a6da061de99dc61b5fb168aedabf62d56d861a18..83da12d2652a3293edcaa819b48bc6004b1c1e03 100644
--- a/paddle/fluid/jit/function_utils.cc
+++ b/paddle/fluid/jit/function_utils.cc
@@ -15,7 +15,9 @@
 #include "paddle/fluid/jit/function_utils.h"
 
 #include "paddle/fluid/framework/program_desc.h"
+#include "paddle/fluid/framework/scope.h"
 #include "paddle/fluid/framework/var_desc.h"
+#include "paddle/fluid/framework/variable.h"
 #include "paddle/phi/core/enforce.h"
 
 namespace paddle {
@@ -75,7 +77,7 @@ void ShareParamsIntoScope(const std::vector<std::string> &param_names,
   for (size_t i = 0; i < param_names.size(); ++i) {
     std::string name = param_names[i];
     auto &param = params_dict.find(name)->second;
-    auto &dense_tensor = param.Get<DenseTensor>();
+    auto &dense_tensor = param->Get<DenseTensor>();
     VLOG(3) << "share into scope: " << name;
     auto *var = scope->Var(name);
     auto *dst_tensor = var->GetMutable<DenseTensor>();
diff --git a/paddle/fluid/jit/function_utils.h b/paddle/fluid/jit/function_utils.h
index ba1eaf7308be9192b1958a2998ed99cad1003b27..90e2e4b7f798f9031fb3d58fd97e1219104f25b2 100644
--- a/paddle/fluid/jit/function_utils.h
+++ b/paddle/fluid/jit/function_utils.h
@@ -18,18 +18,23 @@
 #include <unordered_map>
 #include <vector>
 
-#include "paddle/fluid/framework/scope.h"
-#include "paddle/fluid/framework/variable.h"
 #include "paddle/phi/api/include/tensor.h"
 #include "paddle/phi/common/place.h"
-#include "paddle/phi/core/dense_tensor.h"
 
 #include "paddle/fluid/jit/function_schema.h"
 
 namespace paddle {
+
+namespace framework {
+class Variable;
+class ProgramDesc;
+class Scope;
+}  // namespace framework
+
 namespace jit {
 using Variable = paddle::framework::Variable;
-using Name2VariableMap = std::unordered_map<std::string, Variable>;
+using Name2VariableMap =
+    std::unordered_map<std::string, std::shared_ptr<Variable>>;
 using DenseTensor = phi::DenseTensor;
 using Tensor = paddle::experimental::Tensor;
 
diff --git a/paddle/fluid/jit/layer.cc b/paddle/fluid/jit/layer.cc
index f5985d71b03477d9ddf8b7772754db38e6036978..0e981bc45957f5706df5b32372fb406dd827b6fc 100644
--- a/paddle/fluid/jit/layer.cc
+++ b/paddle/fluid/jit/layer.cc
@@ -14,17 +14,21 @@
 
 #include "paddle/fluid/jit/layer.h"
 
+#include "paddle/fluid/framework/variable.h"
+
+#include "paddle/fluid/jit/base_function.h"
+#include "paddle/fluid/jit/compilation_unit.h"
+#include "paddle/fluid/jit/function_schema.h"
+
 namespace paddle {
 namespace jit {
-Layer::Layer(const std::vector<std::shared_ptr<FunctionInfo>>& infos,
-             const Name2VariableMap& params_dict,
-             const phi::Place& place)
+Layer::Layer(const Name2VariableMap& params_dict, const phi::Place& place)
     : params_dict_(params_dict) {
-  VLOG(3) << "infos size: " << infos.size();
+  unit_.reset(new CompilationUnit());
 }
 
 std::shared_ptr<BaseFunction> Layer::Function(const std::string& name) const {
-  return unit_.Function(name);
+  return unit_->Function(name);
 }
 
 std::vector<Tensor> Layer::forward(const std::vector<Tensor>& inputs) {
@@ -42,15 +46,15 @@ void Layer::to(const phi::Place& place) {}
 
 void Layer::SetFunction(const std::string& name,
                         const std::shared_ptr<BaseFunction>& function) {
-  unit_.SetFunction(name, function);
+  unit_->SetFunction(name, function);
 }
 
 std::vector<std::string> Layer::FunctionNames() const {
-  return unit_.FunctionNames();
+  return unit_->FunctionNames();
 }
 
 const Name2FunctionMap& Layer::FunctionMap() const {
-  return unit_.FunctionMap();
+  return unit_->FunctionMap();
 }
 
 }  // namespace jit
diff --git a/paddle/fluid/jit/layer.h b/paddle/fluid/jit/layer.h
index ee75881fc3156215a921e41c8a8d77bfa5556a07..b2efa77fedf5287f3414a8e20d3d582b6f631544 100644
--- a/paddle/fluid/jit/layer.h
+++ b/paddle/fluid/jit/layer.h
@@ -18,23 +18,31 @@
 #include <unordered_map>
 #include <vector>
 
-#include "paddle/fluid/framework/variable.h"
+#include "paddle/phi/api/include/tensor.h"
 #include "paddle/phi/common/place.h"
 
-#include "paddle/fluid/jit/base_function.h"
-#include "paddle/fluid/jit/compilation_unit.h"
-#include "paddle/fluid/jit/function_schema.h"
+#include "base_function.h"
 
 namespace paddle {
+
+namespace framework {
+class Variable;
+}  // namespace framework
+
 namespace jit {
+class CompilationUnit;
+
+using DenseTensor = phi::DenseTensor;
+using Tensor = paddle::experimental::Tensor;
 using Variable = paddle::framework::Variable;
-using Name2VariableMap = std::unordered_map<std::string, Variable>;
+using Name2VariableMap =
+    std::unordered_map<std::string, std::shared_ptr<Variable>>;
+using Name2FunctionMap =
+    std::unordered_map<std::string, std::shared_ptr<BaseFunction>>;
 
 class Layer {
  public:
-  Layer(const std::vector<std::shared_ptr<FunctionInfo>>& infos,
-        const Name2VariableMap& params_dict,
-        const phi::Place& place);
+  Layer(const Name2VariableMap& params_dict, const phi::Place& place);
 
   std::shared_ptr<BaseFunction> Function(const std::string& name) const;
 
@@ -56,7 +64,7 @@ class Layer {
  private:
   Name2VariableMap params_dict_;
   Name2VariableMap attrs_dict_;
-  CompilationUnit unit_;
+  std::shared_ptr<CompilationUnit> unit_;
 };
 
 }  // namespace jit
diff --git a/paddle/fluid/jit/serializer.cc b/paddle/fluid/jit/serializer.cc
index 2dee9ee879a22de11eaa6a951c0bf45427b7cea2..c24995f71182632edc597062a4eb9d6b1f24f6dd 100644
--- a/paddle/fluid/jit/serializer.cc
+++ b/paddle/fluid/jit/serializer.cc
@@ -16,10 +16,14 @@
 
 #include <set>
 
+#include "paddle/fluid/framework/var_desc.h"
+#include "paddle/fluid/framework/variable.h"
 #include "paddle/fluid/platform/device_context.h"
 
 #include "paddle/fluid/jit/executor_function.h"
+#include "paddle/fluid/jit/layer.h"
 #include "paddle/fluid/jit/pe_function.h"
+#include "paddle/fluid/jit/property.h"
 #include "paddle/fluid/jit/serializer_utils.h"
 
 DECLARE_string(jit_engine_type);
@@ -55,7 +59,7 @@ Layer Deserializer::operator()(const std::string& path,
   ReadTensorData(path + PDPARAMS_SUFFIX, param_names_set, place, &params_dict);
   // ReadAttributeData();
 
-  Layer layer = Layer(infos, params_dict, place);
+  Layer layer = Layer(params_dict, place);
 
   for (auto& info : infos) {
     if (FLAGS_jit_engine_type == "Executor") {
@@ -90,7 +94,7 @@ void Deserializer::ReadTensorData(const std::string& file_name,
     // TODO(dev): Support framework::Vocab
     DenseTensor* dense_tesnor = v.GetMutable<DenseTensor>();
     framework::DeserializeFromStream(fin, dense_tesnor, dev_ctx);
-    (*params_dict)[*it] = v;
+    (*params_dict)[*it] = std::make_shared<Variable>(v);
   }
 }
 
diff --git a/paddle/fluid/jit/serializer.h b/paddle/fluid/jit/serializer.h
index bdc3b81d551938ea3b798c64c3f2e23f6bc6e4b7..188239f469a5768e81436701088a858a47f6753d 100644
--- a/paddle/fluid/jit/serializer.h
+++ b/paddle/fluid/jit/serializer.h
@@ -14,16 +14,26 @@
 
 #pragma once
 
+#include <memory>
+#include <set>
 #include <string>
+#include <unordered_map>
 
-#include "paddle/fluid/framework/var_desc.h"
-#include "paddle/fluid/framework/variable.h"
-#include "paddle/fluid/jit/property.h"
-
-#include "paddle/fluid/jit/layer.h"
+#include "paddle/phi/common/place.h"
 
 namespace paddle {
+
+namespace framework {
+class Variable;
+class ProgramDesc;
+}  // namespace framework
+
 namespace jit {
+class Layer;
+using Variable = paddle::framework::Variable;
+using Name2VariableMap =
+    std::unordered_map<std::string, std::shared_ptr<Variable>>;
+
 // Export Layer into local disk
 class Serializer {
  public:
diff --git a/paddle/fluid/jit/serializer_utils.cc b/paddle/fluid/jit/serializer_utils.cc
index e68d75f58d56d2a1ab4217d961e38a1435ad3dd6..41bfa71b4ce25a1fc61afddca5701fb5ff71b707 100644
--- a/paddle/fluid/jit/serializer_utils.cc
+++ b/paddle/fluid/jit/serializer_utils.cc
@@ -17,6 +17,7 @@
 #include <dirent.h>
 #include <fstream>
 
+#include "paddle/fluid/framework/phi_utils.h"
 #include "paddle/fluid/framework/var_desc.h"
 
 namespace paddle {
@@ -100,6 +101,10 @@ const std::vector<std::pair<std::string, std::string>> PdmodelFilePaths(
   return pdmodel_paths;
 }
 
+void InitKernelSignatureMap() {
+  paddle::framework::InitDefaultKernelSignatureMap();
+}
+
 }  // namespace utils
 }  // namespace jit
 }  // namespace paddle
diff --git a/paddle/fluid/jit/serializer_utils.h b/paddle/fluid/jit/serializer_utils.h
index dfa980544bc31412ff11bee6c99a3fc31ad45aab..97850504d96617cf9b6a23b98512a9ee7edc9202 100644
--- a/paddle/fluid/jit/serializer_utils.h
+++ b/paddle/fluid/jit/serializer_utils.h
@@ -17,9 +17,12 @@
 #include <string>
 #include <vector>
 
-#include "paddle/fluid/framework/var_desc.h"
-
 namespace paddle {
+
+namespace framework {
+class VarDesc;
+}  // namespace framework
+
 namespace jit {
 static const char PDMODEL_SUFFIX[] = ".pdmodel";
 static const char PDPARAMS_SUFFIX[] = ".pdiparams";
@@ -40,6 +43,8 @@ bool FileExists(const std::string& file_path);
 const std::vector<std::pair<std::string, std::string>> PdmodelFilePaths(
     const std::string& path);
 
+void InitKernelSignatureMap();
+
 }  // namespace utils
 }  // namespace jit
 }  // namespace paddle
diff --git a/python/setup.py.in b/python/setup.py.in
index c02ef7f017fca69e20120dccbf0e014e002deab2..1b36b272d0d705fd7772719eff3eb98f47f1a533 100755
--- a/python/setup.py.in
+++ b/python/setup.py.in
@@ -621,8 +621,12 @@ headers = (
     list(find_files('*.h', '@PADDLE_SOURCE_DIR@/paddle/phi/kernels', recursive=True)) +  # phi kernels headers
     # capi headers
     list(find_files('*.h', '@PADDLE_SOURCE_DIR@/paddle/phi/capi', recursive=True)) +  # phi capi headers
-    # utila api headers
-    list(find_files('*.h', '@PADDLE_SOURCE_DIR@/paddle/utils', recursive=True)))  # paddle utils headers
+    # utils api headers
+    list(find_files('*.h', '@PADDLE_SOURCE_DIR@/paddle/utils', recursive=True))) # paddle utils headers
+
+jit_layer_headers = ['layer.h', 'serializer.h', 'serializer_utils.h', 'all.h', 'base_function.h']
+for f in jit_layer_headers:
+    headers += list(find_files(f, '@PADDLE_SOURCE_DIR@/paddle/fluid/jit', recursive=False))
 
 if '${WITH_MKLDNN}' == 'ON':
     headers += list(find_files('*', '${MKLDNN_INSTALL_DIR}/include')) # mkldnn
@@ -667,6 +671,10 @@ class InstallHeaders(Command):
         elif 'third_party' not in header:
             # paddle headers
             install_dir = re.sub('@PADDLE_SOURCE_DIR@/', '', header)
+            print('install_dir: ', install_dir)
+            if 'fluid/jit' in install_dir:
+                install_dir = re.sub('fluid/jit', 'jit', install_dir)
+                print('fluid/jit install_dir: ', install_dir)
         else:
             # third_party
             install_dir = re.sub('${THIRD_PARTY_PATH}', 'third_party', header)