From a3537afc90ba52b0824fef6ef4a87cc72f31478a Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Thu, 30 Jan 2025 14:26:30 -0800 Subject: [PATCH 1/7] Enable Training Pybindings in OSS --- CMakeLists.txt | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index 50575d26a19..9408e2554df 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -791,6 +791,34 @@ if(EXECUTORCH_BUILD_PYBIND) install(TARGETS portable_lib LIBRARY DESTINATION executorch/extension/pybindings ) + + if (EXECUTORCH_BUILD_EXTENSION_TRAINING) + + set(_pybind_training_dep_libs + ${TORCH_PYTHON_LIBRARY} + bundled_program + etdump + executorch + extension_data_loader + util + torch + extension_training + ) + + # pybind training + pybind11_add_module(_training_lib SHARED extension/training/pybindings/_training_lib.cpp) + + target_compile_definitions( + _training_lib PUBLIC EXECUTORCH_PYTHON_MODULE_NAME=_training_lib + ) + 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 extension/training/pybindings + ) + endif() endif() if(EXECUTORCH_BUILD_KERNELS_CUSTOM) From 612e0c12c573d6d2098d8fa6d7709b69c29bfbe7 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Mon, 3 Feb 2025 09:40:22 -0800 Subject: [PATCH 2/7] changes --- CMakeLists.txt | 10 +++++++--- install_executorch.py | 9 ++++++--- setup.py | 16 ++++++++++++++++ 3 files changed, 29 insertions(+), 6 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 9408e2554df..cc336e65761 100644 --- a/CMakeLists.txt +++ b/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() diff --git a/install_executorch.py b/install_executorch.py index 37ef3185ad3..6d111258be2 100644 --- a/install_executorch.py +++ b/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() diff --git a/setup.py b/setup.py index 5e8f155353d..4d61777ce3a 100644 --- a/setup.py +++ b/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( From 32859992a959c447caca60dcc4089cfb75dc6da2 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Mon, 3 Feb 2025 10:27:58 -0800 Subject: [PATCH 3/7] changes --- CMakeLists.txt | 4 +--- install_executorch.py | 3 ++- setup.py | 5 +++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index cc336e65761..8894f79b34f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -812,13 +812,11 @@ if(EXECUTORCH_BUILD_PYBIND) # pybind training pybind11_add_module(_training_lib SHARED extension/training/pybindings/_training_lib.cpp) - target_compile_definitions( - _training_lib PUBLIC EXECUTORCH_PYTHON_MODULE_NAME=_training_lib - ) 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}) + message(STATUS "Installing Training Lib") install(TARGETS _training_lib LIBRARY DESTINATION executorch/extension/training/pybindings ) diff --git a/install_executorch.py b/install_executorch.py index 6d111258be2..fe1e113f193 100644 --- a/install_executorch.py +++ b/install_executorch.py @@ -80,9 +80,10 @@ def main(args): ) if pybind_arg == "training": CMAKE_ARGS += " -DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON" + os.environ["EXECUTORCH_BUILD_TRAINING"] = "ON" else: - EXECUTORCH_BUILD_PYBIND = "ON" CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON" + EXECUTORCH_BUILD_PYBIND = "ON" if args.clean: clean() diff --git a/setup.py b/setup.py index 4d61777ce3a..d00cab83d3e 100644 --- a/setup.py +++ b/setup.py @@ -51,6 +51,7 @@ import platform import re import sys +import logging # Import this before distutils so that setuptools can intercept the distuils # imports. @@ -88,7 +89,7 @@ def pybindings(cls) -> bool: @classmethod def training(cls) -> bool: - return cls._is_env_enabled("EXECUTORCH_BUILD_TRAINING", default=True) + return cls._is_env_enabled("EXECUTORCH_BUILD_TRAINING", default=False) @classmethod def llama_custom_ops(cls) -> bool: @@ -580,6 +581,7 @@ def run(self): "-DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON", ] if ShouldBuild.training(): + logging.info("Building training foobar") cmake_args += [ "-DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON", ] @@ -686,7 +688,6 @@ def get_ext_modules() -> List[Extension]: ) ) if ShouldBuild.training(): - ext_modules.append( # Install the prebuilt pybindings extension wrapper for training BuiltExtension( From 68cec51f2f2daa7438dbc75560467da4e837042d Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Mon, 3 Feb 2025 13:02:38 -0800 Subject: [PATCH 4/7] changes --- CMakeLists.txt | 20 +++++++++++--------- setup.py | 1 + 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 8894f79b34f..8d674ad8010 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -240,6 +240,17 @@ cmake_dependent_option( "NOT EXECUTORCH_BUILD_ARM_BAREMETAL" OFF ) +if(EXECUTORCH_BUILD_PYBIND) + set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER ON) + set(EXECUTORCH_BUILD_DEVTOOLS ON) +endif() + +if(EXECUTORCH_BUILD_EXTENSION_TRAINING) + set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON) + set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER ON) + set(EXECUTORCH_BUILD_EXTENSION_MODULE ON) +endif() + if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT) set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON) set(EXECUTORCH_BUILD_KERNELS_CUSTOM ON) @@ -713,14 +724,6 @@ endif() if(EXECUTORCH_BUILD_PYBIND) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third-party/pybind11) - if(NOT EXECUTORCH_BUILD_EXTENSION_DATA_LOADER) - add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/data_loader) - endif() - - if(NOT EXECUTORCH_BUILD_DEVTOOLS) - add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/devtools) - endif() - # find pytorch lib, to allow pybind to take at::Tensor as input/output find_package(Torch CONFIG REQUIRED) find_library( @@ -816,7 +819,6 @@ if(EXECUTORCH_BUILD_PYBIND) target_compile_options(_training_lib PUBLIC ${_pybind_compile_options}) target_link_libraries(_training_lib PRIVATE ${_pybind_training_dep_libs}) - message(STATUS "Installing Training Lib") install(TARGETS _training_lib LIBRARY DESTINATION executorch/extension/training/pybindings ) diff --git a/setup.py b/setup.py index d00cab83d3e..566bd30d087 100644 --- a/setup.py +++ b/setup.py @@ -585,6 +585,7 @@ def run(self): cmake_args += [ "-DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON", ] + build_args += ["--target", "_training_lib"] 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 From 8e31761ed378691ac0931f80594b8ec9a3006b7f Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Mon, 3 Feb 2025 13:09:14 -0800 Subject: [PATCH 5/7] changes --- setup.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/setup.py b/setup.py index 566bd30d087..fe15dfdd1d4 100644 --- a/setup.py +++ b/setup.py @@ -51,7 +51,6 @@ import platform import re import sys -import logging # Import this before distutils so that setuptools can intercept the distuils # imports. @@ -581,7 +580,6 @@ def run(self): "-DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON", ] if ShouldBuild.training(): - logging.info("Building training foobar") cmake_args += [ "-DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON", ] From 963e8c7f520488a5d28ea04beeacb1684cf2eaf2 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Mon, 3 Feb 2025 13:27:31 -0800 Subject: [PATCH 6/7] lint --- CMakeLists.txt | 2 +- install_executorch.py | 2 +- setup.py | 3 ++- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 8d674ad8010..b7d8cc1e51b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -795,7 +795,7 @@ if(EXECUTORCH_BUILD_PYBIND) LIBRARY DESTINATION executorch/extension/pybindings ) - if (EXECUTORCH_BUILD_EXTENSION_TRAINING) + if(EXECUTORCH_BUILD_EXTENSION_TRAINING) set(_pybind_training_dep_libs ${TORCH_PYTHON_LIBRARY} diff --git a/install_executorch.py b/install_executorch.py index fe1e113f193..9279d9434ea 100644 --- a/install_executorch.py +++ b/install_executorch.py @@ -81,7 +81,7 @@ def main(args): if pybind_arg == "training": CMAKE_ARGS += " -DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON" os.environ["EXECUTORCH_BUILD_TRAINING"] = "ON" - else: + else: CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON" EXECUTORCH_BUILD_PYBIND = "ON" diff --git a/setup.py b/setup.py index fe15dfdd1d4..87c95f0515b 100644 --- a/setup.py +++ b/setup.py @@ -690,7 +690,8 @@ def get_ext_modules() -> List[Extension]: ext_modules.append( # Install the prebuilt pybindings extension wrapper for training BuiltExtension( - "_training_lib.*", "executorch.extension.training.pybindings._training_lib" + "_training_lib.*", + "executorch.extension.training.pybindings._training_lib", ) ) if ShouldBuild.llama_custom_ops(): From 6662f79351f8ceced39881f4643b28f652603273 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Mon, 3 Feb 2025 14:56:30 -0800 Subject: [PATCH 7/7] unbreak weird apple failure --- CMakeLists.txt | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index b7d8cc1e51b..ca8d1bbbcf2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -240,10 +240,6 @@ cmake_dependent_option( "NOT EXECUTORCH_BUILD_ARM_BAREMETAL" OFF ) -if(EXECUTORCH_BUILD_PYBIND) - set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER ON) - set(EXECUTORCH_BUILD_DEVTOOLS ON) -endif() if(EXECUTORCH_BUILD_EXTENSION_TRAINING) set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON) @@ -724,6 +720,14 @@ endif() if(EXECUTORCH_BUILD_PYBIND) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third-party/pybind11) + if(NOT EXECUTORCH_BUILD_EXTENSION_DATA_LOADER) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/data_loader) + endif() + + if(NOT EXECUTORCH_BUILD_DEVTOOLS) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/devtools) + endif() + # find pytorch lib, to allow pybind to take at::Tensor as input/output find_package(Torch CONFIG REQUIRED) find_library(