Change training pybind extension_module linkage to static (#9961)

Summary: Fixes #9576. Use `extension_module_static` in building
`_training_lib`.

Test Plan: Rely on unit test, also did a manual install in editable
mode:

```bash
./install_executorch.sh --pybind training -e
python -c "from executorch.extension.training.pybindings._training_lib import get_sgd_optimizer"
```

Reviewers:

Subscribers:

Tasks:

Tags:

Co-authored-by: Mengwei Liu <[email protected]>

Author: pytorchbot

CMakeLists.txt

@@ -758,10 +758,6 @@ if(EXECUTORCH_BUILD_EXTENSION_MODULE)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/module)
 endif()
 
-if(EXECUTORCH_BUILD_EXTENSION_TRAINING)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/training)
-endif()
-
 if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/runner_util)
 endif()
@@ -862,34 +858,13 @@ if(EXECUTORCH_BUILD_PYBIND)
 
   if(EXECUTORCH_BUILD_EXTENSION_TRAINING)
 
-    set(_pybind_training_dep_libs
-        ${TORCH_PYTHON_LIBRARY}
-        etdump
-        executorch
-        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)
-
-    target_include_directories(_training_lib PRIVATE ${TORCH_INCLUDE_DIRS})
-    target_compile_options(_training_lib PUBLIC ${_pybind_compile_options})
-    target_link_libraries(_training_lib PRIVATE ${_pybind_training_dep_libs})
-
-    install(TARGETS _training_lib
-            LIBRARY DESTINATION executorch/extension/training/pybindings
-    )
   endif()
 endif()
 
+if(EXECUTORCH_BUILD_EXTENSION_TRAINING)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/training)
+endif()
+
 if(EXECUTORCH_BUILD_KERNELS_CUSTOM)
   # TODO: move all custom kernels to ${CMAKE_CURRENT_SOURCE_DIR}/kernels/custom
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/custom_ops)

extension/training/CMakeLists.txt

@@ -26,7 +26,7 @@ target_include_directories(
 target_include_directories(extension_training PUBLIC ${EXECUTORCH_ROOT}/..)
 target_compile_options(extension_training PUBLIC ${_common_compile_options})
 target_link_libraries(extension_training executorch_core
-    extension_data_loader extension_module extension_tensor extension_flat_tensor)
+    extension_data_loader extension_module_static extension_tensor extension_flat_tensor)
 
 
 list(TRANSFORM _train_xor__srcs PREPEND "${EXECUTORCH_ROOT}/")
@@ -40,6 +40,33 @@ train_xor gflags executorch_core portable_ops_lib extension_tensor
 )
 target_compile_options(train_xor PUBLIC ${_common_compile_options})
 
+# Pybind library.
+set(_pybind_training_dep_libs
+  ${TORCH_PYTHON_LIBRARY}
+  etdump
+  executorch
+  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 ${CMAKE_CURRENT_SOURCE_DIR}/pybindings/_training_lib.cpp)
+
+target_include_directories(_training_lib PRIVATE ${TORCH_INCLUDE_DIRS})
+target_compile_options(_training_lib PUBLIC -Wno-deprecated-declarations -fPIC -frtti -fexceptions)
+target_link_libraries(_training_lib PRIVATE ${_pybind_training_dep_libs})
+
+install(TARGETS _training_lib
+    LIBRARY DESTINATION executorch/extension/training/pybindings
+)
+
 # Install libraries
 install(
   TARGETS extension_training

pytest.ini

@@ -51,6 +51,7 @@ addopts =
     extension/llm/modules/test
     extension/llm/export
     extension/pybindings/test
+    extension/training/pybindings/test
     # Runtime
     runtime
     # test TODO: fix these tests

setup.py

@@ -861,7 +861,7 @@ def get_ext_modules() -> List[Extension]:
             ext_modules.append(
                 # Install the prebuilt pybindings extension wrapper for training
                 BuiltExtension(
-                    "_training_lib.*",
+                    "extension/training/_training_lib.*",
                     "executorch.extension.training.pybindings._training_lib",
                 )
             )