fix

Author: flybird11111

colossalai/pipeline/schedule/interleaved_pp.py

@@ -6,7 +6,7 @@
 from torch.nn import Module, ModuleList
 from torch.utils._pytree import tree_map
 
-from colossalai.accelerator import get_accelerator
+from colossalai.accelerator import get_accelerator, BaseAccelerator
 from colossalai.interface import OptimizerWrapper
 from colossalai.pipeline.p2p import PipelineP2PCommunication, create_send_metadata
 from colossalai.pipeline.stage_manager import PipelineStageManager
@@ -17,7 +17,7 @@
 from .base import PipelineSchedule
 
 
-def _wait_p2p(wait_handles: List[get_accelerator().Event]) -> None:
+def _wait_p2p(wait_handles: List[BaseAccelerator.Event]) -> None:
     if wait_handles is not None:
         for req in wait_handles:
             req.wait()

colossalai/shardformer/layer/normalization.py

@@ -17,11 +17,9 @@
 SUPPORT_NPU = False
 try:
     import torch_npu
-
     SUPPORT_NPU = True
-    warnings.warn("support npu")
 except Exception:
-    warnings.warn("support gpu")
+    pass
 
 
 __all__ = ["FusedLayerNorm", "FusedRMSNorm", "LayerNorm", "RMSNorm", "BaseLayerNorm"]

examples/language/llama/benchmark.py

@@ -340,7 +340,7 @@ def empty_init():
 
     torch.set_default_dtype(torch.float)
     coordinator.print_on_master(
-        f"Booster init max NPU memory: {get_accelerator().max_memory_allocated()/1024**2:.2f} MB"
+        f"Booster init max device memory: {get_accelerator().max_memory_allocated()/1024**2:.2f} MB"
     )
     coordinator.print_on_master(
         f"Booster init max CPU memory: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024:.2f} MB"

extensions/pybind/flash_attention/flash_attention_npu.py

@@ -1,4 +1,5 @@
 from ...base_extension import _Extension
+import math
 
 
 class FlashAttentionNpuExtension(_Extension):
@@ -27,7 +28,6 @@ def build_jit(self) -> None:
         )
 
     def load(self):
-        import math
         from typing import Optional
 
         import torch