[REDUX] Refactor Apex build process to use the PyTorch JIT extension flow #291
Conversation
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
…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
….all_ops BuilderUtils class
…e readable names.
… 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.
… tests in the contrib test runner
…o run jit build and tests in readme, add other tests in readme
…files in apex folder
…on or CUDAExtension is built with JIT load approach
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.