NVIDIA / apex

A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch
BSD 3-Clause "New" or "Revised" License
8.42k stars 1.4k forks source link

Option to only build `amp_C` module #1725

Open ezhang887 opened 1 year ago

ezhang887 commented 1 year ago

Hi!

Currently to build the amp_C module I need to specify --cuda_ext.

However with --cuda_ext the build can be pretty slow sometimes, and a lot of the time is spent building the other modules.

If I only want to use stuff from the amp_C module, having an option to build just that would make the build much faster than --cuda_ext which builds 7 more modules in addition to amp_C.

Maybe cc @crcrpar based on recent activity?