gitextract_8yaiblk9/ ├── .clang-format ├── .git-blame-ignore-revs ├── .github/ │ └── ISSUE_TEMPLATE/ │ └── bug_report.md ├── .gitignore ├── .gitmodules ├── .nojekyll ├── .pre-commit-config.yaml ├── LICENSE ├── README.md ├── apex/ │ ├── __init__.py │ ├── _autocast_utils.py │ ├── contrib/ │ │ ├── __init__.py │ │ ├── bottleneck/ │ │ │ ├── __init__.py │ │ │ ├── bottleneck.py │ │ │ ├── halo_exchangers.py │ │ │ └── test.py │ │ ├── clip_grad/ │ │ │ ├── __init__.py │ │ │ └── clip_grad.py │ │ ├── conv_bias_relu/ │ │ │ ├── __init__.py │ │ │ └── conv_bias_relu.py │ │ ├── csrc/ │ │ │ ├── bottleneck/ │ │ │ │ └── bottleneck.cpp │ │ │ ├── conv_bias_relu/ │ │ │ │ └── conv_bias_relu.cpp │ │ │ ├── cudnn_gbn/ │ │ │ │ ├── cudnn_gbn.cpp │ │ │ │ ├── norm_sample.cpp │ │ │ │ └── norm_sample.h │ │ │ ├── fmha/ │ │ │ │ ├── fmha_api.cpp │ │ │ │ └── src/ │ │ │ │ ├── fmha/ │ │ │ │ │ ├── gemm.h │ │ │ │ │ ├── gmem_tile.h │ │ │ │ │ ├── kernel_traits.h │ │ │ │ │ ├── mask.h │ │ │ │ │ ├── smem_tile.h │ │ │ │ │ ├── softmax.h │ │ │ │ │ └── utils.h │ │ │ │ ├── fmha.h │ │ │ │ ├── fmha_dgrad_fp16_128_64_kernel.sm80.cu │ │ │ │ ├── fmha_dgrad_fp16_256_64_kernel.sm80.cu │ │ │ │ ├── fmha_dgrad_fp16_384_64_kernel.sm80.cu │ │ │ │ ├── fmha_dgrad_fp16_512_64_kernel.sm80.cu │ │ │ │ ├── fmha_dgrad_kernel_1xN_reload.h │ │ │ │ ├── fmha_dgrad_kernel_1xN_reload_nl.h │ │ │ │ ├── fmha_fill.cu │ │ │ │ ├── fmha_fprop_fp16_128_64_kernel.sm80.cu │ │ │ │ ├── fmha_fprop_fp16_256_64_kernel.sm80.cu │ │ │ │ ├── fmha_fprop_fp16_384_64_kernel.sm80.cu │ │ │ │ ├── fmha_fprop_fp16_512_64_kernel.sm80.cu │ │ │ │ ├── fmha_fprop_kernel_1xN.h │ │ │ │ ├── fmha_kernel.h │ │ │ │ ├── fmha_noloop_reduce.cu │ │ │ │ └── fmha_utils.h │ │ │ ├── focal_loss/ │ │ │ │ ├── focal_loss_cuda.cpp │ │ │ │ └── focal_loss_cuda_kernel.cu │ │ │ ├── gpu_direct_storage/ │ │ │ │ ├── gds.cpp │ │ │ │ ├── gds.h │ │ │ │ └── gds_pybind.cpp │ │ │ ├── group_norm/ │ │ │ │ ├── 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 │ │ │ ├── group_norm_v2/ │ │ │ │ ├── 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 │ │ │ ├── groupbn/ │ │ │ │ ├── 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/ │ │ │ │ ├── index_mul_2d_cuda.cpp │ │ │ │ └── index_mul_2d_cuda_kernel.cu │ │ │ ├── layer_norm/ │ │ │ │ ├── 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 │ │ │ ├── multihead_attn/ │ │ │ │ ├── additive_masked_softmax_dropout_cuda.cu │ │ │ │ ├── dropout.cuh │ │ │ │ ├── encdec_multihead_attn_cuda.cu │ │ │ │ ├── encdec_multihead_attn_norm_add_cuda.cu │ │ │ │ ├── layer_norm.cuh │ │ │ │ ├── masked_softmax_dropout_cuda.cu │ │ │ │ ├── multihead_attn_frontend.cpp │ │ │ │ ├── philox.cuh │ │ │ │ ├── self_multihead_attn_bias_additive_mask_cuda.cu │ │ │ │ ├── self_multihead_attn_bias_cuda.cu │ │ │ │ ├── self_multihead_attn_cuda.cu │ │ │ │ ├── self_multihead_attn_norm_add_cuda.cu │ │ │ │ ├── softmax.cuh │ │ │ │ └── strided_batched_gemm.cuh │ │ │ ├── nccl_allocator/ │ │ │ │ └── NCCLAllocator.cpp │ │ │ ├── nccl_p2p/ │ │ │ │ ├── nccl_p2p.cpp │ │ │ │ ├── nccl_p2p_cuda.cu │ │ │ │ ├── nccl_p2p_cuda.cuh │ │ │ │ ├── nccl_version.cpp │ │ │ │ └── nccl_version_check.cu │ │ │ ├── optimizers/ │ │ │ │ ├── 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/ │ │ │ │ ├── peer_memory.cpp │ │ │ │ ├── peer_memory_cuda.cu │ │ │ │ └── peer_memory_cuda.cuh │ │ │ ├── transducer/ │ │ │ │ ├── transducer_joint.cpp │ │ │ │ ├── transducer_joint_kernel.cu │ │ │ │ ├── transducer_loss.cpp │ │ │ │ └── transducer_loss_kernel.cu │ │ │ └── xentropy/ │ │ │ ├── interface.cpp │ │ │ └── xentropy_kernel.cu │ │ ├── cudnn_gbn/ │ │ │ ├── __init__.py │ │ │ └── batch_norm.py │ │ ├── examples/ │ │ │ ├── gpu_direct_storage/ │ │ │ │ ├── benchmark_load.py │ │ │ │ ├── benchmark_save.py │ │ │ │ ├── example_load.py │ │ │ │ └── example_save.py │ │ │ ├── multihead_attn/ │ │ │ │ ├── func_test_multihead_attn.py │ │ │ │ └── perf_test_multihead_attn.py │ │ │ └── nccl_allocator/ │ │ │ ├── allreduce.py │ │ │ ├── cache.py │ │ │ ├── change_cuda_allocator.py │ │ │ └── toy_ddp.py │ │ ├── fmha/ │ │ │ ├── __init__.py │ │ │ └── fmha.py │ │ ├── focal_loss/ │ │ │ ├── __init__.py │ │ │ └── focal_loss.py │ │ ├── gpu_direct_storage/ │ │ │ ├── README.md │ │ │ └── __init__.py │ │ ├── group_norm/ │ │ │ ├── __init__.py │ │ │ └── group_norm.py │ │ ├── groupbn/ │ │ │ ├── __init__.py │ │ │ └── batch_norm.py │ │ ├── index_mul_2d/ │ │ │ ├── __init__.py │ │ │ └── index_mul_2d.py │ │ ├── layer_norm/ │ │ │ ├── __init__.py │ │ │ └── layer_norm.py │ │ ├── multihead_attn/ │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ ├── encdec_multihead_attn.py │ │ │ ├── encdec_multihead_attn_func.py │ │ │ ├── fast_encdec_multihead_attn_func.py │ │ │ ├── fast_encdec_multihead_attn_norm_add_func.py │ │ │ ├── fast_self_multihead_attn_func.py │ │ │ ├── fast_self_multihead_attn_norm_add_func.py │ │ │ ├── mask_softmax_dropout_func.py │ │ │ ├── self_multihead_attn.py │ │ │ └── self_multihead_attn_func.py │ │ ├── nccl_allocator/ │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ └── nccl_allocator.py │ │ ├── openfold_triton/ │ │ │ ├── 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 │ │ ├── optimizers/ │ │ │ ├── __init__.py │ │ │ ├── distributed_fused_adam.py │ │ │ ├── distributed_fused_lamb.py │ │ │ ├── fp16_optimizer.py │ │ │ ├── fused_adam.py │ │ │ ├── fused_lamb.py │ │ │ └── fused_sgd.py │ │ ├── peer_memory/ │ │ │ ├── __init__.py │ │ │ ├── peer_halo_exchanger_1d.py │ │ │ └── peer_memory.py │ │ ├── sparsity/ │ │ │ ├── COPYRIGHT │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ ├── asp.py │ │ │ ├── permutation_lib.py │ │ │ ├── permutation_search_kernels/ │ │ │ │ ├── CUDA_kernels/ │ │ │ │ │ └── permutation_search_kernels.cu │ │ │ │ ├── __init__.py │ │ │ │ ├── call_permutation_search_kernels.py │ │ │ │ ├── channel_swap.py │ │ │ │ ├── exhaustive_search.py │ │ │ │ └── permutation_utilities.py │ │ │ ├── permutation_tests/ │ │ │ │ ├── README.md │ │ │ │ ├── ablation_studies.sh │ │ │ │ ├── permutation_test.py │ │ │ │ ├── runtime_table.sh │ │ │ │ └── unstructured_study.sh │ │ │ ├── sparse_masklib.py │ │ │ └── test/ │ │ │ ├── checkpointing_test_part1.py │ │ │ ├── checkpointing_test_part2.py │ │ │ ├── checkpointing_test_reference.py │ │ │ ├── test_permutation_application.py │ │ │ └── toy_problem.py │ │ ├── test/ │ │ │ ├── __init__.py │ │ │ ├── bottleneck/ │ │ │ │ ├── __init__.py │ │ │ │ └── test_bottleneck_module.py │ │ │ ├── clip_grad/ │ │ │ │ ├── __init__.py │ │ │ │ └── test_clip_grad.py │ │ │ ├── conv_bias_relu/ │ │ │ │ ├── __init__.py │ │ │ │ └── test_conv_bias_relu.py │ │ │ ├── cudnn_gbn/ │ │ │ │ ├── __init__.py │ │ │ │ └── test_cudnn_gbn_with_two_gpus.py │ │ │ ├── fmha/ │ │ │ │ ├── __init__.py │ │ │ │ └── test_fmha.py │ │ │ ├── focal_loss/ │ │ │ │ ├── __init__.py │ │ │ │ └── test_focal_loss.py │ │ │ ├── fused_dense/ │ │ │ │ └── test_fused_dense.py │ │ │ ├── group_norm/ │ │ │ │ ├── __init__.py │ │ │ │ └── test_group_norm.py │ │ │ ├── index_mul_2d/ │ │ │ │ ├── __init__.py │ │ │ │ └── test_index_mul_2d.py │ │ │ ├── layer_norm/ │ │ │ │ ├── __init__.py │ │ │ │ └── test_fast_layer_norm.py │ │ │ ├── multihead_attn/ │ │ │ │ ├── __init__.py │ │ │ │ ├── test_encdec_multihead_attn.py │ │ │ │ ├── test_encdec_multihead_attn_norm_add.py │ │ │ │ ├── test_fast_self_multihead_attn_bias.py │ │ │ │ ├── test_mha_fused_softmax.py │ │ │ │ ├── test_self_multihead_attn.py │ │ │ │ └── test_self_multihead_attn_norm_add.py │ │ │ ├── openfold_triton/ │ │ │ │ ├── test_fused_adam_swa.py │ │ │ │ ├── test_openfold_mha.py │ │ │ │ └── test_sync_triton_auto_tune_cache_across_gpus.py │ │ │ ├── optimizers/ │ │ │ │ ├── __init__.py │ │ │ │ ├── test_dist_adam.py │ │ │ │ └── test_distributed_fused_lamb.py │ │ │ ├── peer_memory/ │ │ │ │ ├── __init__.py │ │ │ │ └── test_peer_halo_exchange_module.py │ │ │ ├── transducer/ │ │ │ │ ├── __init__.py │ │ │ │ ├── test_transducer_joint.py │ │ │ │ └── test_transducer_loss.py │ │ │ └── xentropy/ │ │ │ ├── __init__.py │ │ │ └── test_label_smoothing.py │ │ ├── torchsched/ │ │ │ ├── __init__.py │ │ │ ├── backend.py │ │ │ ├── config.py │ │ │ ├── inductor/ │ │ │ │ ├── __init__.py │ │ │ │ ├── _utils.py │ │ │ │ ├── event.py │ │ │ │ ├── graph.py │ │ │ │ ├── scheduler.py │ │ │ │ └── wrapper.py │ │ │ ├── ops/ │ │ │ │ ├── __init__.py │ │ │ │ └── layer_norm.py │ │ │ └── passes/ │ │ │ ├── __init__.py │ │ │ └── pre_grad_passes.py │ │ ├── transducer/ │ │ │ ├── __init__.py │ │ │ ├── _transducer_ref.py │ │ │ └── transducer.py │ │ └── xentropy/ │ │ ├── __init__.py │ │ └── softmax_xentropy.py │ ├── distributed_testing/ │ │ ├── __init__.py │ │ ├── _ucc_util.py │ │ └── distributed_test_base.py │ ├── fused_dense/ │ │ ├── __init__.py │ │ └── fused_dense.py │ ├── mlp/ │ │ ├── __init__.py │ │ └── mlp.py │ ├── multi_tensor_apply/ │ │ ├── __init__.py │ │ └── multi_tensor_apply.py │ ├── normalization/ │ │ ├── __init__.py │ │ └── fused_layer_norm.py │ └── optimizers/ │ ├── __init__.py │ ├── fused_adagrad.py │ ├── fused_adam.py │ ├── fused_lamb.py │ ├── fused_mixed_precision_lamb.py │ ├── fused_novograd.py │ └── fused_sgd.py ├── csrc/ │ ├── amp_C_frontend.cpp │ ├── flatten_unflatten.cpp │ ├── fused_dense.cpp │ ├── fused_dense_cuda.cu │ ├── layer_norm_cuda.cpp │ ├── layer_norm_cuda_kernel.cu │ ├── megatron/ │ │ ├── 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 │ ├── 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 ├── docs/ │ ├── Makefile │ └── source/ │ ├── _static/ │ │ └── css/ │ │ └── pytorch_theme.css │ ├── _templates/ │ │ └── layout.html │ ├── conf.py │ ├── index.rst │ ├── layernorm.rst │ └── optimizers.rst ├── examples/ │ ├── README.md │ ├── dcgan/ │ │ ├── README.md │ │ └── main_amp.py │ ├── docker/ │ │ ├── Dockerfile │ │ └── README.md │ ├── imagenet/ │ │ ├── README.md │ │ └── main_amp.py │ └── simple/ │ └── distributed/ │ ├── README.md │ ├── distributed_data_parallel.py │ └── run.sh ├── pyproject.toml ├── requirements.txt ├── requirements_dev.txt ├── setup.py └── tests/ ├── L0/ │ ├── run_fused_layer_norm/ │ │ └── test_fused_layer_norm.py │ ├── run_mlp/ │ │ └── test_mlp.py │ ├── run_optimizers/ │ │ ├── __init__.py │ │ ├── test_adam.py │ │ ├── test_fused_novograd.py │ │ ├── test_fused_optimizer.py │ │ └── test_lamb.py │ └── run_test.py ├── L1/ │ ├── common/ │ │ ├── compare.py │ │ ├── main_amp.py │ │ └── run_test.sh │ ├── cross_product/ │ │ └── run.sh │ └── cross_product_distributed/ │ └── run.sh ├── distributed/ │ ├── DDP/ │ │ ├── ddp_race_condition_test.py │ │ └── run_race_test.sh │ ├── amp_master_params/ │ │ ├── amp_master_params.py │ │ ├── compare.py │ │ └── run.sh │ └── synced_batchnorm/ │ ├── 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 └── docker_extension_builds/ └── run.sh