drop use_kernels_from_apex

Author: mayank31398

megatron/arguments.py

@@ -641,8 +641,6 @@ def _add_training_args(parser):
                        'training if SIGTERM is received')
     group.add_argument('--tensorboard-dir', type=str, default=None,
                        help='Write TensorBoard logs to this directory.')
-    group.add_argument("--use-kernels-from-apex", action="store_true",
-                       help="use Apex kernels instead of Megatron")
     group.add_argument('--no-masked-softmax-fusion',
                        action='store_false',
                        help='Disable fusion of query_key_value scaling, '

megatron/fused_kernels/__init__.py

@@ -2,9 +2,6 @@
 
 
 def load(args):
-    if args.use_kernels_from_apex:
-        return
-
     if torch.version.hip is None:
         print("running on CUDA devices")
         from megatron.fused_kernels.cuda import load as load_kernels

megatron/fused_kernels/tests/test_fused_kernels.py

@@ -283,7 +283,6 @@ class DummyArgs:
     rank: int = 0
     masked_softmax_fusion: bool = True
     gradient_accumulation_fusion: bool = True
-    use_kernels_from_apex: bool = False
 
 
 if __name__ == "__main__":

tools/checkpoint_loader_megatron.py

@@ -5,7 +5,6 @@
 
 import torch
 
-
 def add_arguments(parser):
     group = parser.add_argument_group(title='Megatron loader')
 
@@ -16,8 +15,6 @@ def add_arguments(parser):
                        'trim padding from the embedding table.')
     group.add_argument('--megatron-path', type=str, default=None,
                        help='Base directory of deepspeed repository')
-    group.add_argument("--use-kernels-from-apex", action="store_true",
-                       help="use Apex kernels instead of Megatron")
 
 def _load_checkpoint(queue, args):
 

tools/checkpoint_saver_megatron.py

@@ -6,7 +6,6 @@
 
 import torch
 
-
 def add_arguments(parser):
     group = parser.add_argument_group(title='Megatron saver')
 
@@ -19,8 +18,6 @@ def add_arguments(parser):
     group.add_argument('--target-pipeline-parallel-size', type=int,
                        help='Target tensor model parallel size, default to the pipeline parall size '
                        'in the input checkpoint if provided by the loader, otherwise to 1')
-    group.add_argument("--use-kernels-from-apex", action="store_true",
-                       help="use Apex kernels instead of Megatron")
 
 def save_checkpoint(queue, args):