[REDUX] Refactor Apex build process to use the PyTorch JIT extension flow#291
Merged
jithunnair-amd merged 83 commits intomasterfrom Jan 26, 2026
Merged
[REDUX] Refactor Apex build process to use the PyTorch JIT extension flow#291jithunnair-amd merged 83 commits intomasterfrom
jithunnair-amd merged 83 commits intomasterfrom
Conversation
…ad of building it. Code uses accelerator and op_builder modules from deepspeed code.
…lly created by setup.py for the build process
… jit mode, add csrc back to setup.py since it is not copied to apex wheel
… building the wheel
… imports in python module
… make MLP JIT compile
…piled during apex installation
…thod to CUDAOpBuilder to support its jit compile
…daOpBuilder support jit of this module
… nvcc_threads_args method in CUDAOpBuilder to support these jit modules
…as it is needed for TorchCPUOpBuilder
…CUDAOpBuilder to support this
…repo is cloned in the docker
…across the list entries
Collaborator
|
! cherry-pick --onto release/1.10.0 |
|
Can't perform the cherry-pick keyword: unexpected error Comment processed by Build |
amd-sriram
added a commit
that referenced
this pull request
Feb 4, 2026
…flow (#291) * Created initial code for loading fused_dense module dynamically instead of building it. Code uses accelerator and op_builder modules from deepspeed code. * add apex/git_version_info_installed.py to gitignore as it is dynamically created by setup.py for the build process * add code for building fused rope dynamically * add code for building fused bias swiglu dynamically * fix the code so that fused rope and fused softmax are not compiled in jit mode, add csrc back to setup.py since it is not copied to apex wheel * load the jit modules inside and this prevents them from building when building the wheel * convert syncbn module to jit * fix the unnecessary compile of syncbn module in wheel building due to imports in python module * add fused layer norm module to jit build * make focal loss module as jit module * make focal loss module as jit module * make xentropy module as jit module * make bpn module as jit module * add code to build individual extensions without JIT * clean up the flags for the modules based on apex/setup.py * add function to get the backward_pass_guard_args in CudaOpBuilder and make MLP JIT compile * add fused weight gradient mlp to jit compile * move fused_weight_gradient_mlp_cuda load inside so that it is not compiled during apex installation * make fused index mul 2d jit compile and dd aten atomic header flag method to CUDAOpBuilder to support its jit compile * make fast multihead attention as jit module, add generator_args to CudaOpBuilder support jit of this module * make transducer loss and transducer joint modules as jit modules, add nvcc_threads_args method in CUDAOpBuilder to support these jit modules * remove extra method - installed_cuda_version from CUDAOpBuilder * add apex_C module to jit compile, add py-cpuinfo to requirements.txt as it is needed for TorchCPUOpBuilder * make nccl allocator as a jit compile module, add nccl_args method to CUDAOpBuilder to support this * make amp_C as a jit module * add a few uses of amp_C jit module * add a few uses of amp_C jit module * make fused adam as a jit module * add a few uses of amp_C jit module * fix the issue with fused adam jit module * make fused lamb as jit module * make distributed adam as jit module * make distributed lamb as jit module * add remaining amp_C uses with jit loader * add remaining usage of apexC jit module * make nccl p2p module as jit compile * make peer memory module as jit compile * add code to check for minimum nccl version to compile nccl allocator module * add provision to provide APEX_CPP_OPS=1 and APEX_CUDA_OPS=1 as replacement for --cpp_ext --cuda_ext command line arguments for building specific extensions in apex, save these settings for later use * check for minimum torch version for nccl allocator, check if the module is compatible other removed from installed ops list * add build as a dependency to support wheel building * Replace is_compatible to check for installation conditions with is_supported, because there is an issue with loading nccl allocator * Similar to pytorch we create a make command to install aiter, that the user can use. There will be no building aiter in the setup.py * update extension import test so that it considers jit compile extensions * clean up MultiTensorApply usages so that amp_C is not build in jit compile mode * Adding missing modules from deepspeed repo. Remove extra code in setup.py. Use is_compatible instead of is_supported * change name of apex_C module * change the name of cpp and cuda build flags, remove APEX_BUILD_OPS, cleanup the logic to build specific modules * add missing files used in cpu accelerator * add make clean command to handle deleting torch extensions installed for jit modules, fix the cpu builder import error * remove unused code in setup.py, fix the code to build for cpu mode * Removing unused code * remove accelerator package and refactor the used code into op_builder.all_ops BuilderUtils class * remove accelerator package usages * revert code that was removed by mistake * Cleaning up the setup file and renaming functions and variable to more readable names. * Fix the nccl version so that the nccl_allocator.so file can be loaded properly. Setup() call has an argument called py_modules which copies the python class into sitepackages folder. The python modules in the compatibility folder do lazy load of the builder classes. First these files are copied in the parent folder so that the files themselves are copied into sitepackages so that the kernel can be loaded into python then these temporary files are deleted. * Restore to original importing the extension code. * renamed compatibility/scaled_masked_softmax_cuda.py, added some extra tests in the contrib test runner * Added instructions for JIT load and changes in installation options * Restructuring the README * Added instructions for building wheel * replaced TorchCPUBuilder with CPUBuilder, added a main method in contrib test runner * create a script to build different jit conditions for running different tests * add script to run tests with different jit builds, add instructions to run jit build and tests in readme, add other tests in readme * fix the issues with running the tests - improper paths, counting .so files in apex folder * add mad internal scripts * remove print statement * remove testing section from readme * change location of result file * remove multiple results file from models.json * add platform specific description to wheel name even if no CppExtension or CUDAExtension is built with JIT load approach * add ninja and wheel to requirements to be installed * Update Release notes in Readme * Exclude compatibility folder while installing apex * Update README.md * Update README.md * Update README.md * Adding modification note to the original copywrite * fix the issue with symbolic links for op_builder, csrc when the apex repo is cloned in the docker * assign the symbolically linked folders into a variable and then loop across the list entries * remove unnecessary tabs --------- Co-authored-by: skishore <sriramkumar.kishorekumar@amd.com> Co-authored-by: sriram <sriram.kumar@silo.ai>
amd-sriram
added a commit
that referenced
this pull request
Feb 4, 2026
…flow (#291) (#296) * Created initial code for loading fused_dense module dynamically instead of building it. Code uses accelerator and op_builder modules from deepspeed code. * add apex/git_version_info_installed.py to gitignore as it is dynamically created by setup.py for the build process * add code for building fused rope dynamically * add code for building fused bias swiglu dynamically * fix the code so that fused rope and fused softmax are not compiled in jit mode, add csrc back to setup.py since it is not copied to apex wheel * load the jit modules inside and this prevents them from building when building the wheel * convert syncbn module to jit * fix the unnecessary compile of syncbn module in wheel building due to imports in python module * add fused layer norm module to jit build * make focal loss module as jit module * make focal loss module as jit module * make xentropy module as jit module * make bpn module as jit module * add code to build individual extensions without JIT * clean up the flags for the modules based on apex/setup.py * add function to get the backward_pass_guard_args in CudaOpBuilder and make MLP JIT compile * add fused weight gradient mlp to jit compile * move fused_weight_gradient_mlp_cuda load inside so that it is not compiled during apex installation * make fused index mul 2d jit compile and dd aten atomic header flag method to CUDAOpBuilder to support its jit compile * make fast multihead attention as jit module, add generator_args to CudaOpBuilder support jit of this module * make transducer loss and transducer joint modules as jit modules, add nvcc_threads_args method in CUDAOpBuilder to support these jit modules * remove extra method - installed_cuda_version from CUDAOpBuilder * add apex_C module to jit compile, add py-cpuinfo to requirements.txt as it is needed for TorchCPUOpBuilder * make nccl allocator as a jit compile module, add nccl_args method to CUDAOpBuilder to support this * make amp_C as a jit module * add a few uses of amp_C jit module * add a few uses of amp_C jit module * make fused adam as a jit module * add a few uses of amp_C jit module * fix the issue with fused adam jit module * make fused lamb as jit module * make distributed adam as jit module * make distributed lamb as jit module * add remaining amp_C uses with jit loader * add remaining usage of apexC jit module * make nccl p2p module as jit compile * make peer memory module as jit compile * add code to check for minimum nccl version to compile nccl allocator module * add provision to provide APEX_CPP_OPS=1 and APEX_CUDA_OPS=1 as replacement for --cpp_ext --cuda_ext command line arguments for building specific extensions in apex, save these settings for later use * check for minimum torch version for nccl allocator, check if the module is compatible other removed from installed ops list * add build as a dependency to support wheel building * Replace is_compatible to check for installation conditions with is_supported, because there is an issue with loading nccl allocator * Similar to pytorch we create a make command to install aiter, that the user can use. There will be no building aiter in the setup.py * update extension import test so that it considers jit compile extensions * clean up MultiTensorApply usages so that amp_C is not build in jit compile mode * Adding missing modules from deepspeed repo. Remove extra code in setup.py. Use is_compatible instead of is_supported * change name of apex_C module * change the name of cpp and cuda build flags, remove APEX_BUILD_OPS, cleanup the logic to build specific modules * add missing files used in cpu accelerator * add make clean command to handle deleting torch extensions installed for jit modules, fix the cpu builder import error * remove unused code in setup.py, fix the code to build for cpu mode * Removing unused code * remove accelerator package and refactor the used code into op_builder.all_ops BuilderUtils class * remove accelerator package usages * revert code that was removed by mistake * Cleaning up the setup file and renaming functions and variable to more readable names. * Fix the nccl version so that the nccl_allocator.so file can be loaded properly. Setup() call has an argument called py_modules which copies the python class into sitepackages folder. The python modules in the compatibility folder do lazy load of the builder classes. First these files are copied in the parent folder so that the files themselves are copied into sitepackages so that the kernel can be loaded into python then these temporary files are deleted. * Restore to original importing the extension code. * renamed compatibility/scaled_masked_softmax_cuda.py, added some extra tests in the contrib test runner * Added instructions for JIT load and changes in installation options * Restructuring the README * Added instructions for building wheel * replaced TorchCPUBuilder with CPUBuilder, added a main method in contrib test runner * create a script to build different jit conditions for running different tests * add script to run tests with different jit builds, add instructions to run jit build and tests in readme, add other tests in readme * fix the issues with running the tests - improper paths, counting .so files in apex folder * add mad internal scripts * remove print statement * remove testing section from readme * change location of result file * remove multiple results file from models.json * add platform specific description to wheel name even if no CppExtension or CUDAExtension is built with JIT load approach * add ninja and wheel to requirements to be installed * Update Release notes in Readme * Exclude compatibility folder while installing apex * Update README.md * Update README.md * Update README.md * Adding modification note to the original copywrite * fix the issue with symbolic links for op_builder, csrc when the apex repo is cloned in the docker * assign the symbolically linked folders into a variable and then loop across the list entries * remove unnecessary tabs --------- Co-authored-by: Jithun Nair <37884920+jithunnair-amd@users.noreply.github.com> Co-authored-by: sriram <sriram.kumar@silo.ai>
jithunnair-amd
pushed a commit
to ROCm/pytorch
that referenced
this pull request
Feb 4, 2026
[REDUX] Refactor Apex build process to use the PyTorch JIT extension flow ([#291](ROCm/apex#291)) ([#296](ROCm/apex#296))
Collaborator
|
! cherry-pick --onto release/1.9.0 |
Collaborator
|
! cherry-pick --onto release/1.8.0 |
|
Created branch autogenerated/release/1.9.0_cherry-pick_pr-291 and #305. It contains a merge conflict. Please resolve it Comment processed by Build |
|
Created branch autogenerated/release/1.8.0_cherry-pick_pr-291 and #306. It contains a merge conflict. Please resolve it Comment processed by Build |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Had to revert #247 due to a build breakage seen in AISW HUD runs. Removed the change from
masterbranch until we can figure out the root cause. This PR will be used to re-merge the changes.Motivation
Currently building apex takes around 30 minutes. The motivation behind JIT (just in time) load is to reduce the installation time to under 1 minute and then build the modules on demand when using them or running their tests.
In addition, this PR provides flexibility in building specific modules (based on the argument to indicate CPP or CUDA extension or specific modules).
Due to https://discuss.python.org/t/symbolic-links-in-wheels/1945/19, we process the symbolic links and remove them by copying the contents where they pointed to the location of symbolic links before we call setup in setup.py. The changes were tested locally as well as on AISW HUD runs.
Technical Details
To install from source:
pip install . --no-build-isolationTo build the wheel and install from wheel:
Currently converted extensions include (which work on rocm)
Total - 27 extensions
The following extensions have been not been included in jit load in this PR as these are not used before (Nvidia gpu only)
Added custom code to support building apex modules
Other changes
make aitersimilar to pytorch.make cleanto remove torch extensions created with JIT load.Tested Unit tests
cd tests/L0
PYTHONUNBUFFERED=1 sh run_rocm.sh 2>&1 | tee log_results.txt
cd apex/contrib/test/
PYTHONUNBUFFERED=1 python3 run_rocm_extensions.py 2>&1 | tee log_results_contrib.txt
torchrun --nproc_per_node 8 apex/contrib/peer_memory/peer_halo_exchange_module_tests.py
cd tests/distributed/synced_batchnorm
sh unit_test.sh
Docker used for the testing (tested with CPU only and with GPU)
registry-sc-harbor.amd.com/framework/compute-rocm-rel-7.0:32_ubuntu22.04_py3.10_pytorch_release-2.8_d2d97084
Tested following commands
Running the different build instructions.
madengine run --tags pyt_deepspeed_megatron_llama2_7b --live-output --additional-context "{'guest_os': 'UBUNTU', 'docker_build_arg':{'BASE_DOCKER':'registry-sc-harbor.amd.com/framework/compute-rocm-dkms-no-npi-hipclang:16771_ubuntu24.04_py3.12_pytorch_rocm7.1_internal_testing_d1fb13a8'}}"Running extensive tests on MI300
Created an issue for the two errors: https://github.com/ROCm/frameworks-internal/issues/14438
Running a few commands on MI200
Creating the wheels in a docker with CPU and running the tests with GPU docker (MI300)
Created scripts for testing JIT build and documented the instructions at https://amd.atlassian.net/wiki/spaces/MLSE/pages/1255652200/Testing