fused

Author: mayank31398

megatron/fused_kernels/__init__.py

@@ -6,12 +6,10 @@ def load(args):
         return
 
     if torch.version.hip is None:
-        if torch.distributed.get_rank() == 0:
-            print("running on CUDA devices")
+        print("running on CUDA devices")
         from megatron.fused_kernels.cuda import load as load_kernels
     else:
-        if torch.distributed.get_rank() == 0:
-            print("running on ROCm devices")
+        print("running on ROCm devices")
         from megatron.fused_kernels.rocm import load as load_kernels
 
     load_kernels(args)

megatron/initialize.py

@@ -31,7 +31,7 @@
 except ModuleNotFoundError:
     print('Wandb import failed', flush=True)
 
-import megatron.fused_kernels as fused_kernels
+from megatron import fused_kernels
 from megatron import get_adlr_autoresume
 from megatron import get_args
 from megatron import get_tensorboard_writer

tools/checkpoint_loader_megatron.py

@@ -5,7 +5,7 @@
 
 import torch
 
-import megatron.fused_kernels as fused_kernels
+from megatron import fused_kernels
 
 def add_arguments(parser):
     group = parser.add_argument_group(title='Megatron loader')

tools/checkpoint_saver_megatron.py

@@ -6,7 +6,7 @@
 
 import torch
 
-import megatron.fused_kernels as fused_kernels
+from megatron import fused_kernels
 
 def add_arguments(parser):
     group = parser.add_argument_group(title='Megatron saver')