[Ascend] Native Integration of MegatronAdaptor (FlagScale Module) on FlagOS

flagscale/train/megatron/plugin_flagscale/dist_signal_handler.py

@@ -3,6 +3,7 @@
 import torch
 
 from megatron.plugin.platform import get_platform
+from functools import wraps
 cur_platform = get_platform()
 
 logger = logging.getLogger(__name__)

flagscale/train/megatron/plugin_flagscale/npu_plugin.py

@@ -0,0 +1,37 @@
+import time
+
+import torch
+import torch_npu
+
+def get_device_arch_version():
+    return 8
+
+def _compile_dependencies():
+    if torch.distributed.get_rank() == 0:
+        start_time = time.time()
+        print('> compiling dataset index builder ...')
+        from megatron.core.datasets.utils import compile_helpers
+
+        compile_helpers()
+        print(
+            '>>> done with dataset index builder. Compilation time: {:.3f} seconds'.format(time.time() - start_time),
+            flush=True,
+        )
+
+def get_device(local_rank=None):
+    backend = torch.distributed.get_backend()
+    if backend == 'hccl':
+        if local_rank is None:
+            device = torch.device('cuda')
+        else:
+            device = torch.device(f'cuda:{local_rank}')
+    elif backend == 'nccl':
+        if local_rank is None:
+            device = torch.device(cur_platform.device_name())
+        else:
+            device = torch.device(f'{cur_platform.device_name()}:{local_rank}')
+    elif backend == 'gloo':
+        device = torch.device('cpu')
+    else:
+        raise RuntimeError(f"Unsupported distributed backend: {backend}")
+    return device

flagscale/train/megatron/plugin_flagscale/override_registry.py

@@ -8,7 +8,26 @@
 # =============================================================================
 # DistSignalHandler - get_device
 # =============================================================================
+
 register(
     target="megatron.training.dist_signal_handler.get_device",
     impl="megatron.plugin_flagscale.dist_signal_handler.get_device",
 )
+
+register(
+    target="megatron.training.dist_signal_handler.get_device",
+    impl="megatron.plugin_flagscale.npu_plugin.get_device",
+    vendor="npu",
+)
+
+register(
+    target="megatron.training.utils.get_device_arch_version",
+    impl="megatron.plugin_flagscale.npu_plugin.get_device_arch_version",
+    vendor="npu",
+)
+
+register(
+    target="megatron.training.initialize.initialize._compile_dependencies",
+    impl="megatron.plugin_flagscale.npu_plugin._compile_dependencies",
+    vendor="npu",
+)

flagscale/train/megatron/training/initialize.py

@@ -37,6 +37,7 @@
 from megatron.training.arguments_fs import FSTrainArguments
 from megatron.training.global_vars import set_spiky_loss_detector
 from megatron.plugin.hetero.parallel_context import set_parallel_context
+from megatron.plugin.decorators import overridable
 
 logger = logging.getLogger(__name__)
 
@@ -198,7 +199,7 @@ def finish_mpu_init():
         # No continuation function
         return None
 
-
+@overridable
 def _compile_dependencies():
 
     args = get_args()

flagscale/train/megatron/training/utils.py

@@ -48,6 +48,7 @@
 from megatron.plugin.utils import get_device_type_for_comm
 from megatron.plugin.platform import get_platform
 cur_platform = get_platform()
+from megatron.plugin.decorators import overridable
 
 def calc_params_l2_norm(model, force_create_fp32_copy=False):
     """Calculate l2 norm of parameters"""
@@ -543,6 +544,7 @@ def is_first_or_last_pipeline_stage(vp_stage, dualpipev_stage=None):
     )
 
 
+@overridable
 def get_device_arch_version():
     """Returns GPU arch version (8: Ampere, 9: Hopper, 10: Blackwell, ...)"""
     return cur_platform.get_device_properties(cur_platform.device(0)).major