fix

Author: mayank31398

megatron/fused_kernels/__init__.py

@@ -1,10 +1,17 @@
+import torch
+
+
 def load(args):
     if args.use_kernels_from_apex:
         return
 
-    if args.device == "cuda":
+    if torch.version.hip is None:
+        if torch.distributed.get_rank() == 0:
+            print("running on CUDA devices")
         from megatron.fused_kernels.cuda import load as load_kernels
-    elif args.device == "rocm":
+    else:
+        if torch.distributed.get_rank() == 0:
+            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)
 
-from megatron.fused_kernels import cuda
+import megatron.fused_kernels as fused_kernels
 from megatron import get_adlr_autoresume
 from megatron import get_args
 from megatron import get_tensorboard_writer
@@ -198,11 +198,11 @@ def _compile_dependencies():
     if torch.distributed.get_rank() == 0:
         start_time = time.time()
         print('> compiling and loading fused kernels ...', flush=True)
-        cuda.load(args)
+        fused_kernels.load(args)
         torch.distributed.barrier()
     else:
         torch.distributed.barrier()
-        cuda.load(args)
+        fused_kernels.load(args)
     # Simple barrier to make sure all ranks have passed the
     # compilation phase successfully before moving on to the
     # rest of the program. We think this might ensure that

tools/checkpoint_loader_megatron.py

@@ -5,7 +5,7 @@
 
 import torch
 
-from megatron.fused_kernels import cuda
+import megatron.fused_kernels as fused_kernels
 
 def add_arguments(parser):
     group = parser.add_argument_group(title='Megatron loader')
@@ -133,7 +133,7 @@ def get_models(count, dtype, pre_process, post_process):
     set_global_variables(margs)
     mpu.initialize.set_tensor_model_parallel_world_size(margs.tensor_model_parallel_size)
     mpu.initialize.set_pipeline_model_parallel_world_size(margs.pipeline_model_parallel_size)
-    cuda.load(margs)
+    fused_kernels.load(margs)
 
     # Get true (non-padded) vocab size
     if args.true_vocab_size is not None:

tools/checkpoint_saver_megatron.py

@@ -6,7 +6,7 @@
 
 import torch
 
-from megatron.fused_kernels import cuda
+import megatron.fused_kernels as fused_kernels
 
 def add_arguments(parser):
     group = parser.add_argument_group(title='Megatron saver')
@@ -161,7 +161,7 @@ def get_models(count, dtype, pre_process, post_process):
     mpu.initialize.set_pipeline_model_parallel_world_size(args.target_pipeline_parallel_size)
     mpu.initialize.set_tensor_model_parallel_rank(0)
     mpu.initialize.set_pipeline_model_parallel_rank(0)
-    cuda.load(margs)
+    fused_kernels.load(margs)
 
     # Embeddings
     #-----------