ljx adapt npu

Author: li126com

internlm/core/context/parallel_context.py

@@ -11,7 +11,7 @@
 import torch
 import torch.distributed as dist
 
-from internlm.accelerator import get_accelerator
+from internlm.accelerator import AcceleratorType, get_accelerator
 from internlm.utils.common import SingletonMeta
 from internlm.utils.config import Config
 from internlm.utils.logger import get_logger
@@ -309,7 +309,10 @@ def init_global_dist(self, rank: int, world_size: int, backend: str, host: str,
            use_cpu (bool): whether to set up cpu process group.
         """
         # initialize the default process group
-        init_method = f"tcp://[{host}]:{port}"
+        if internlm_accelerator.get_accelerator_backend() == AcceleratorType.GPU:
+            init_method = f"tcp://[{host}]:{port}"
+        else:
+            init_method = f"tcp://{host}:{port}"
         dist.init_process_group(
             rank=rank,
             world_size=world_size,

internlm/core/engine.py

@@ -3,23 +3,31 @@
 
 # adopted from https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/engine
 
-from contextlib import nullcontext
 from typing import List, Optional
 
 import torch
-import transformer_engine.pytorch as te
 from torch.nn import Module
 from torch.nn.modules.loss import _Loss
 from torch.optim.lr_scheduler import _LRScheduler
-from transformer_engine.common.recipe import DelayedScaling, Format
 
+from internlm.accelerator import AcceleratorType, get_accelerator
 from internlm.core.context import ParallelMode
 from internlm.core.context import global_context as gpc
 from internlm.core.gradient_handler import BaseGradientHandler
 from internlm.solver.optimizer import BaseOptimizer
 from internlm.solver.schedulers import Beta2Scheduler
 from internlm.utils.common import get_batch_size, move_to_device
 
+try:
+    from contextlib import nullcontext
+
+    import transformer_engine.pytorch as te
+    from transformer_engine.common.recipe import DelayedScaling, Format
+except ImportError:
+    pass
+
+internlm_accelerator = get_accelerator()
+
 
 class Engine:
     """
@@ -83,27 +91,28 @@ def __init__(
         # build gradient handler
         self._gradient_handlers = gradient_handlers if gradient_handlers else []
 
-        # FP8 GEMM
-        fp8_cfg = gpc.config.get("fp8", None)
-        self.use_fp8 = fp8_cfg is not None
-        self.fp8_recipe = None
-        self.fp8_group = None
-        if self.use_fp8:
-            self.fp8_group = gpc.get_group(ParallelMode.GLOBAL)
-            if fp8_cfg.format == "e4m3":
-                fp8_format = Format.E4M3
-            elif fp8_cfg.format == "hybrid":
-                fp8_format = Format.HYBRID
-            else:
-                raise ValueError("The DelayedScaling recipe only supports E4M3 and HYBRID formats.")
-            self.fp8_recipe = DelayedScaling(
-                margin=fp8_cfg.margin,
-                interval=fp8_cfg.interval,
-                fp8_format=fp8_format,
-                amax_history_len=fp8_cfg.amax_history_len,
-                amax_compute_algo=fp8_cfg.amax_compute_algo,
-                override_linear_precision=(False, False, not fp8_cfg.fp8_wgrad),
-            )
+        if internlm_accelerator.get_accelerator_backend() == AcceleratorType.GPU:
+            # FP8 GEMM
+            fp8_cfg = gpc.config.get("fp8", None)
+            self.use_fp8 = fp8_cfg is not None
+            self.fp8_recipe = None
+            self.fp8_group = None
+            if self.use_fp8:
+                self.fp8_group = gpc.get_group(ParallelMode.GLOBAL)
+                if fp8_cfg.format == "e4m3":
+                    fp8_format = Format.E4M3
+                elif fp8_cfg.format == "hybrid":
+                    fp8_format = Format.HYBRID
+                else:
+                    raise ValueError("The DelayedScaling recipe only supports E4M3 and HYBRID formats.")
+                self.fp8_recipe = DelayedScaling(
+                    margin=fp8_cfg.margin,
+                    interval=fp8_cfg.interval,
+                    fp8_format=fp8_format,
+                    amax_history_len=fp8_cfg.amax_history_len,
+                    amax_compute_algo=fp8_cfg.amax_compute_algo,
+                    override_linear_precision=(False, False, not fp8_cfg.fp8_wgrad),
+                )
 
     @property
     def model(self):
@@ -193,11 +202,13 @@ def __call__(self, *args, **kwargs):
         Returns:
             torch.Tensor: The output of the model.
         """
-        with te.fp8_autocast(
-            enabled=self.use_fp8, fp8_recipe=self.fp8_recipe, fp8_group=self.fp8_group
-        ) if self.use_fp8 else nullcontext():
-            output = self.model(*args, **kwargs)
-        return output
+        if internlm_accelerator.get_accelerator_backend() == AcceleratorType.GPU:
+            with te.fp8_autocast(
+                enabled=self.use_fp8, fp8_recipe=self.fp8_recipe, fp8_group=self.fp8_group
+            ) if self.use_fp8 else nullcontext():
+                output = self.model(*args, **kwargs)
+            return output
+        return self.model(*args, **kwargs)
 
     def load_batch(self, data_iter, to_gpu=True):
         """

internlm/core/trainer_builder.py

@@ -8,6 +8,7 @@
 import torch.distributed as dist
 from torch.utils.data import DataLoader
 
+from internlm.accelerator import AcceleratorType, get_accelerator
 from internlm.checkpoint.checkpoint_manager import CheckpointManager
 from internlm.core.context import ParallelMode
 from internlm.core.context import global_context as gpc
@@ -48,6 +49,7 @@
 
 # global llm logger
 logger = logging.getLogger(__file__)
+internlm_accelerator = get_accelerator()
 
 
 class TrainerBuilder(Trainer):
@@ -114,7 +116,8 @@ def __init__(
         criterion = self._initialize_criterion()
 
         # initialize cpu offload manager for selective checkpoint
-        initialize_offload_manager(gpc.config.get("selective_checkpoint_offload", False))
+        if internlm_accelerator.get_accelerator_backend() == AcceleratorType.GPU:
+            initialize_offload_manager(gpc.config.get("selective_checkpoint_offload", False))
 
         # initialize train state
         train_state = get_train_state(train_dl)

internlm/model/model_implementations/builder.py

@@ -1,9 +1,9 @@
 from typing import List, Union
 
 import torch
-import transformer_engine.pytorch as te
 from torch import nn
 
+from internlm.accelerator import AcceleratorType, get_accelerator
 from internlm.core.context import ParallelMode
 from internlm.core.context import global_context as gpc
 from internlm.core.parallel.shard import pipeline_parallel_sharding_wrapper
@@ -20,7 +20,13 @@
 from internlm.utils.logger import get_logger
 from internlm.utils.parallel import is_using_fsdp, is_using_hf, is_using_isp
 
+try:
+    import transformer_engine.pytorch as te
+except ImportError:
+    pass
+
 logger = get_logger(__file__)
+internlm_accelerator = get_accelerator()
 
 
 def simple_swap(model, device):
@@ -156,7 +162,8 @@ def traverse(module):
         else:
             traverse(model)
 
-    if gpc.config.get("fp8", None) is not None:
-        simple_swap(model, fsdp_init_method)
+    if internlm_accelerator.get_accelerator_backend() == AcceleratorType.GPU:
+        if gpc.config.get("fp8", None) is not None:
+            simple_swap(model, fsdp_init_method)
 
     return model

setup.py

@@ -27,7 +27,7 @@ def get_requires() -> List[str]:
 
 extra_require = {
     "torch": ["torch>=2.1.0"],
-    "torch-npu": ["torch==2.1.0", "torch-npu==2.1.0.post3"],
+    "torch-npu": ["torch==2.1.0", "torch-npu==2.1.0.post3", "numpy==1.26.4", "scipy", "decorator"],
 }
 
 setup(