changes

Author: JacobSzwejbka

CMakeLists.txt

@@ -796,15 +796,19 @@ if(EXECUTORCH_BUILD_PYBIND)
 
     set(_pybind_training_dep_libs
         ${TORCH_PYTHON_LIBRARY}
-        bundled_program
         etdump
         executorch
-        extension_data_loader
         util
         torch
         extension_training
     )
 
+    if(EXECUTORCH_BUILD_XNNPACK)
+      # need to explicitly specify XNNPACK and microkernels-prod
+      # here otherwise uses XNNPACK and microkernel-prod symbols from libtorch_cpu
+      list(APPEND _pybind_training_dep_libs xnnpack_backend XNNPACK microkernels-prod)
+    endif()
+
     # pybind training
     pybind11_add_module(_training_lib SHARED extension/training/pybindings/_training_lib.cpp)
 
@@ -816,7 +820,7 @@ if(EXECUTORCH_BUILD_PYBIND)
     target_link_libraries(_training_lib PRIVATE ${_pybind_training_dep_libs})
 
     install(TARGETS _training_lib
-            LIBRARY DESTINATION extension/training/pybindings
+            LIBRARY DESTINATION executorch/extension/training/pybindings
     )
   endif()
 endif()

install_executorch.py

@@ -32,7 +32,7 @@ def clean():
     print("Done cleaning build artifacts.")
 
 
-VALID_PYBINDS = ["coreml", "mps", "xnnpack"]
+VALID_PYBINDS = ["coreml", "mps", "xnnpack", "training"]
 
 
 def main(args):
@@ -78,8 +78,11 @@ def main(args):
                     raise Exception(
                         f"Unrecognized pybind argument {pybind_arg}; valid options are: {', '.join(VALID_PYBINDS)}"
                     )
-                EXECUTORCH_BUILD_PYBIND = "ON"
-                CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON"
+                if pybind_arg == "training":
+                    CMAKE_ARGS += " -DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON"
+                else: 
+                    EXECUTORCH_BUILD_PYBIND = "ON"
+                    CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON"
 
     if args.clean:
         clean()

setup.py

@@ -86,6 +86,10 @@ def _is_env_enabled(env_var: str, default: bool = False) -> bool:
     def pybindings(cls) -> bool:
         return cls._is_env_enabled("EXECUTORCH_BUILD_PYBIND", default=False)
 
+    @classmethod
+    def training(cls) -> bool:
+        return cls._is_env_enabled("EXECUTORCH_BUILD_TRAINING", default=True)
+
     @classmethod
     def llama_custom_ops(cls) -> bool:
         return cls._is_env_enabled("EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT", default=True)
@@ -575,6 +579,10 @@ def run(self):
                 "-DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON",  # add quantized ops to pybindings.
                 "-DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON",
             ]
+            if ShouldBuild.training():
+                cmake_args += [
+                    "-DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON",
+                ]
             build_args += ["--target", "portable_lib"]
             # To link backends into the portable_lib target, callers should
             # add entries like `-DEXECUTORCH_BUILD_XNNPACK=ON` to the CMAKE_ARGS
@@ -677,6 +685,14 @@ def get_ext_modules() -> List[Extension]:
                 "_portable_lib.*", "executorch.extension.pybindings._portable_lib"
             )
         )
+        if ShouldBuild.training():
+
+            ext_modules.append(
+                # Install the prebuilt pybindings extension wrapper for training
+                BuiltExtension(
+                    "_training_lib.*", "executorch.extension.training.pybindings._training_lib"
+                )
+            )
     if ShouldBuild.llama_custom_ops():
         ext_modules.append(
             BuiltFile(