apex
A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch
File Explorer
Download Latest Version (.zip)- bug_report.md
- __init__.py
- bottleneck.py
- halo_exchangers.py
- test.py
- __init__.py
- clip_grad.py
- __init__.py
- conv_bias_relu.py
- bottleneck.cpp
- conv_bias_relu.cpp
- cudnn_gbn.cpp
- norm_sample.cpp
- norm_sample.h
- focal_loss_cuda.cpp
- focal_loss_cuda_kernel.cu
- group_norm_nhwc.cpp
- group_norm_nhwc.h
- group_norm_nhwc_bwd_one_pass.h
- group_norm_nhwc_bwd_one_pass_kernel.cuh
- group_norm_nhwc_bwd_two_pass.cu
- group_norm_nhwc_fwd_one_pass.h
- group_norm_nhwc_fwd_one_pass_kernel.cuh
- group_norm_nhwc_fwd_two_pass.cu
- group_norm_nhwc_one_pass_10.cu
- group_norm_nhwc_one_pass_112.cu
- group_norm_nhwc_one_pass_12.cu
- group_norm_nhwc_one_pass_120.cu
- group_norm_nhwc_one_pass_128.cu
- group_norm_nhwc_one_pass_14.cu
- group_norm_nhwc_one_pass_16.cu
- group_norm_nhwc_one_pass_160.cu
- group_norm_nhwc_one_pass_20.cu
- group_norm_nhwc_one_pass_24.cu
- group_norm_nhwc_one_pass_26.cu
- group_norm_nhwc_one_pass_28.cu
- group_norm_nhwc_one_pass_30.cu
- group_norm_nhwc_one_pass_32.cu
- group_norm_nhwc_one_pass_4.cu
- group_norm_nhwc_one_pass_40.cu
- group_norm_nhwc_one_pass_42.cu
- group_norm_nhwc_one_pass_48.cu
- group_norm_nhwc_one_pass_56.cu
- group_norm_nhwc_one_pass_60.cu
- group_norm_nhwc_one_pass_64.cu
- group_norm_nhwc_one_pass_70.cu
- group_norm_nhwc_one_pass_8.cu
- group_norm_nhwc_one_pass_80.cu
- group_norm_nhwc_one_pass_84.cu
- group_norm_nhwc_one_pass_96.cu
- group_norm_nhwc_one_pass_98.cu
- group_norm_nhwc_op.cpp
- macros.h
- traits.h
- generate_gn_cuda_inst.py
- gn.cpp
- gn.hpp
- gn_cuda.cu
- gn_cuda_host_template.cuh
- gn_cuda_inst_1024_1280.cu
- gn_cuda_inst_1024_1920.cu
- gn_cuda_inst_1024_320.cu
- gn_cuda_inst_1024_640.cu
- gn_cuda_inst_1024_960.cu
- gn_cuda_inst_256_1280.cu
- gn_cuda_inst_256_1920.cu
- gn_cuda_inst_256_2560.cu
- gn_cuda_inst_256_640.cu
- gn_cuda_inst_4096_320.cu
- gn_cuda_inst_4096_640.cu
- gn_cuda_inst_4096_960.cu
- gn_cuda_inst_64_1280.cu
- gn_cuda_inst_64_2560.cu
- gn_cuda_kernel.cuh
- gn_dispatch_hw_c.hpp
- gn_utils.cpp
- gn_utils.hpp
- batch_norm.cu
- batch_norm.h
- batch_norm_add_relu.cu
- batch_norm_add_relu.h
- cuda_utils.h
- interface.cpp
- ipc.cu
- nhwc_batch_norm_kernel.h
- index_mul_2d_cuda.cpp
- index_mul_2d_cuda_kernel.cu
- ln.h
- ln_api.cpp
- ln_bwd_kernels.cuh
- ln_bwd_semi_cuda_kernel.cu
- ln_fwd_cuda_kernel.cu
- ln_fwd_kernels.cuh
- ln_kernel_traits.h
- ln_utils.cuh
- NCCLAllocator.cpp
- nccl_p2p.cpp
- nccl_p2p_cuda.cu
- nccl_p2p_cuda.cuh
- nccl_version.cpp
- nccl_version_check.cu
- fused_adam_cuda.cpp
- fused_adam_cuda_kernel.cu
- fused_lamb_cuda.cpp
- fused_lamb_cuda_kernel.cu
- multi_tensor_distopt_adam.cpp
- multi_tensor_distopt_adam_kernel.cu
- multi_tensor_distopt_lamb.cpp
- multi_tensor_distopt_lamb_kernel.cu
- peer_memory.cpp
- peer_memory_cuda.cu
- peer_memory_cuda.cuh
- philox.cuh
- transducer_joint.cpp
- transducer_joint_kernel.cu
- transducer_loss.cpp
- transducer_loss_kernel.cu
- interface.cpp
- xentropy_kernel.cu
- cudnn-frontend
- __init__.py
- batch_norm.py
- benchmark_load.py
- benchmark_save.py
- example_load.py
- example_save.py
- allreduce.py
- cache.py
- change_cuda_allocator.py
- toy_ddp.py
- __init__.py
- focal_loss.py
- __init__.py
- group_norm.py
- __init__.py
- batch_norm.py
- __init__.py
- index_mul_2d.py
- __init__.py
- layer_norm.py
- __init__.py
- nccl_allocator.py
- README.md
- __init__.py
- _layer_norm_backward_kernels.py
- _layer_norm_config_ampere.py
- _layer_norm_config_hopper.py
- _layer_norm_forward_kernels.py
- _mha_kernel.py
- fused_adam_swa.py
- layer_norm.py
- mha.py
- README.md
- __init__.py
- distributed_fused_adam.py
- distributed_fused_lamb.py
- fp16_optimizer.py
- fused_adam.py
- fused_lamb.py
- fused_sgd.py
- __init__.py
- peer_halo_exchanger_1d.py
- peer_memory.py
- permutation_search_kernels.cu
- __init__.py
- call_permutation_search_kernels.py
- channel_swap.py
- exhaustive_search.py
- permutation_utilities.py
- ablation_studies.sh
- permutation_test.py
- README.md
- runtime_table.sh
- unstructured_study.sh
- checkpointing_test_part1.py
- checkpointing_test_part2.py
- checkpointing_test_reference.py
- test_permutation_application.py
- toy_problem.py
- __init__.py
- asp.py
- COPYRIGHT
- permutation_lib.py
- README.md
- sparse_masklib.py
- __init__.py
- test_bottleneck_module.py
- __init__.py
- test_clip_grad.py
- __init__.py
- test_conv_bias_relu.py
- __init__.py
- test_cudnn_gbn_with_two_gpus.py
- __init__.py
- test_focal_loss.py
- test_fused_dense.py
- __init__.py
- test_group_norm.py
- __init__.py
- test_index_mul_2d.py
- __init__.py
- test_fast_layer_norm.py
- test_fused_adam_swa.py
- test_openfold_mha.py
- test_sync_triton_auto_tune_cache_across_gpus.py
- __init__.py
- test_dist_adam.py
- test_distributed_fused_lamb.py
- __init__.py
- test_peer_halo_exchange_module.py
- __init__.py
- test_transducer_joint.py
- test_transducer_loss.py
- __init__.py
- test_label_smoothing.py
- __init__.py
- __init__.py
- _utils.py
- event.py
- graph.py
- scheduler.py
- wrapper.py
- __init__.py
- layer_norm.py
- __init__.py
- pre_grad_passes.py
- __init__.py
- backend.py
- config.py
- __init__.py
- _transducer_ref.py
- transducer.py
- __init__.py
- softmax_xentropy.py
- __init__.py
- __init__.py
- _ucc_util.py
- distributed_test_base.py
- __init__.py
- fused_dense.py
- __init__.py
- mlp.py
- __init__.py
- multi_tensor_apply.py
- __init__.py
- fused_layer_norm.py
- __init__.py
- fused_adagrad.py
- fused_adam.py
- fused_lamb.py
- fused_mixed_precision_lamb.py
- fused_novograd.py
- fused_sgd.py
- __init__.py
- _autocast_utils.py
- fused_rotary_positional_embedding.cpp
- fused_rotary_positional_embedding.h
- fused_rotary_positional_embedding_cuda.cu
- fused_weight_gradient_dense.cpp
- fused_weight_gradient_dense_16bit_prec_cuda.cu
- fused_weight_gradient_dense_cuda.cu
- generic_scaled_masked_softmax.cpp
- generic_scaled_masked_softmax.h
- generic_scaled_masked_softmax_cuda.cu
- scaled_masked_softmax.cpp
- scaled_masked_softmax.h
- scaled_masked_softmax_cuda.cu
- scaled_softmax.cpp
- scaled_softmax_cuda.cu
- scaled_upper_triang_masked_softmax.cpp
- scaled_upper_triang_masked_softmax.h
- scaled_upper_triang_masked_softmax_cuda.cu
- amp_C_frontend.cpp
- flatten_unflatten.cpp
- fused_dense.cpp
- fused_dense_cuda.cu
- layer_norm_cuda.cpp
- layer_norm_cuda_kernel.cu
- mlp.cpp
- mlp_cuda.cu
- multi_tensor_adagrad.cu
- multi_tensor_adam.cu
- multi_tensor_apply.cuh
- multi_tensor_axpby_kernel.cu
- multi_tensor_l2norm_kernel.cu
- multi_tensor_l2norm_kernel_mp.cu
- multi_tensor_l2norm_scale_kernel.cu
- multi_tensor_lamb.cu
- multi_tensor_lamb_mp.cu
- multi_tensor_lamb_stage_1.cu
- multi_tensor_lamb_stage_2.cu
- multi_tensor_novograd.cu
- multi_tensor_scale_kernel.cu
- multi_tensor_sgd_kernel.cu
- static_switch.h
- syncbn.cpp
- type_shim.h
- update_scale_hysteresis.cu
- welford.cu
- pytorch_theme.css
- nv-pytorch2.png
- layout.html
- conf.py
- index.rst
- layernorm.rst
- optimizers.rst
- Makefile
- main_amp.py
- README.md
- Dockerfile
- README.md
- main_amp.py
- README.md
- distributed_data_parallel.py
- README.md
- run.sh
- README.md
- amp_master_params.py
- compare.py
- run.sh
- ddp_race_condition_test.py
- run_race_test.sh
- python_single_gpu_unit_test.py
- single_gpu_unit_test.py
- test_batchnorm1d.py
- test_groups.py
- two_gpu_test_different_batch_size.py
- two_gpu_unit_test.py
- unit_test.sh
- run.sh
- test_fused_layer_norm.py
- test_mlp.py
- __init__.py
- test_adam.py
- test_fused_novograd.py
- test_fused_optimizer.py
- test_lamb.py
- run_test.py
- compare.py
- main_amp.py
- run_test.sh
- run.sh
- run.sh
- .clang-format
- .git-blame-ignore-revs
- .gitignore
- .gitmodules
- .nojekyll
- .pre-commit-config.yaml
- LICENSE
- pyproject.toml
- README.md
- requirements.txt
- requirements_dev.txt
- setup.py
๐ Installation Guide
1. Get the code
git clone https://github.com/NVIDIA/apex
Downloads the entire project code from GitHub to your computer.
cd apex
Moves into the project folder you just downloaded.
2. Official Install Script
Easy RecommendedPrerequisites
- Python 3 Python is required to use pip.
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
Installs the package published on PyPI directly โ no need to clone the source.
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./
Installs the package published on PyPI directly โ no need to clone the source.
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--group_norm" ./
Installs the package published on PyPI directly โ no need to clone the source.
pip install -v --disable-pip-version-check --no-build-isolation --no-cache-dir ./
Installs the package published on PyPI directly โ no need to clone the source.
After installing, open a new terminal and run the program's version command (e.g. --version) to confirm it worked.
Pulled directly from this repo's README.
3. Docker
EasyPrerequisites
- Git Needed to download the project code from GitHub.
- Docker Desktop Needed to build and run containers. Install it and keep it running in the background.
docker build -f examples/docker/Dockerfile -t apex .
Builds a runnable image based on the Dockerfile.
docker run -p 8080:80 apex
Runs the built image as an actual container.
Run docker compose ps to check the containers are Up. If the README mentions a port, open http://localhost:PORT in your browser.
4. Python
EasyPrerequisites
APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .
Installs the Python libraries listed in requirements.txt (or similar).
APEX_CPP_EXT=1 APEX_CUDA_EXT=1 APEX_FUSED_CONV_BIAS_RELU=1 pip install -v --no-build-isolation .
Installs the Python libraries listed in requirements.txt (or similar).
APEX_CPP_EXT=1 APEX_CUDA_EXT=1 APEX_ALL_CONTRIB_EXT=1 pip install -v --no-build-isolation .
Installs the Python libraries listed in requirements.txt (or similar).
NVCC_APPEND_FLAGS="--threads 4" APEX_PARALLEL_BUILD=8 APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .
Installs the Python libraries listed in requirements.txt (or similar).
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
Installs the package published on PyPI directly โ no need to clone the source.
If it runs without errors and prints output in the terminal, it worked.
Pulled directly from this repo's README.
5. Make
MediumPrerequisites
- Git Needed to download the project code from GitHub.
- Make Usually pre-installed on Linux/macOS. On Windows, install separately (e.g. via MSYS2 or WSL).
cd docs
This project's files live in a subfolder, so move into it first.
make
Compiles the code based on the generated build configuration to produce an executable.
If it finishes without errors, it worked. Try running the generated executable directly.
// repository documentation
Was this content helpful?
(0 ratings)