diff --git a/mindnlp/core/_tensor.py b/mindnlp/core/_tensor.py index 7ee9ea8f7..c122e0b2d 100644 --- a/mindnlp/core/_tensor.py +++ b/mindnlp/core/_tensor.py @@ -812,6 +812,9 @@ def record_stream(self, stream): Tensor.record_stream = record_stream StubTensor.record_stream = record_stream + Tensor.scatter = ops.scatter + StubTensor.scatter = ops.scatter + def _rebuild_from_type_v2(func, new_type, args, state): ret = func(*args) return ret \ No newline at end of file diff --git a/mindnlp/core/configs.py b/mindnlp/core/configs.py index 6543b55c1..4711dfcc8 100644 --- a/mindnlp/core/configs.py +++ b/mindnlp/core/configs.py @@ -4,7 +4,7 @@ SOC = MSContext.get_instance().get_ascend_soc_version() DEVICE_TARGET = mindspore.get_context('device_target') -SUPPORT_BF16 = DEVICE_TARGET == 'Ascend' and SOC not in ['ascend910', 'ascend310b1', 'ascend310b4'] +SUPPORT_BF16 = DEVICE_TARGET == 'Ascend' and SOC not in ['ascend910', 'ascend310b'] ON_A1 = not SUPPORT_BF16 ON_ORANGE_PI = '310b' in SOC USE_PYBOOST = DEVICE_TARGET == 'Ascend' diff --git a/mindnlp/core/npu/__init__.py b/mindnlp/core/npu/__init__.py index c5ded5b1f..febdda27c 100644 --- a/mindnlp/core/npu/__init__.py +++ b/mindnlp/core/npu/__init__.py @@ -2,9 +2,14 @@ import mindspore from mindspore import get_rng_state, set_rng_state, manual_seed -from mindspore.hal import * +from mindspore.runtime import memory_reserved as ms_memory_reserved, \ + memory_allocated as ms_memory_allocated, StreamCtx as StreamContext, Stream, empty_cache, \ + reset_peak_memory_stats, reset_max_memory_allocated, max_memory_allocated, synchronize, \ + current_stream +from mindspore.device_context.ascend import device_count from mindnlp import core +from ..configs import SUPPORT_BF16 FloatTensor = core.FloatTensor HalfTensor = core.FloatTensor @@ -28,6 +33,15 @@ def set_device(device): def _lazy_call(callable, **kwargs): callable() +def is_bf16_supported(): + return SUPPORT_BF16 + +def memory_allocated(device=None): + return ms_memory_allocated() + +def memory_reserved(device=None): + return ms_memory_reserved() + class device: r"""Context-manager that changes the selected device. diff --git a/mindnlp/core/ops/array.py b/mindnlp/core/ops/array.py index d1ed73f28..f01a9ab90 100644 --- a/mindnlp/core/ops/array.py +++ b/mindnlp/core/ops/array.py @@ -71,11 +71,13 @@ def gather(input, dim, index): _complex = _get_cache_prim(ops.Complex)() return _complex(real_part, imag_part) - if use_pyboost() and has_gather: + if use_pyboost() and has_gather and not ON_ORANGE_PI: return mindspore.mint.gather(input, dim, index) - index = ops.where(index < input.shape[dim], index, index - input.shape[dim]) - return ops.gather_elements(input, dim, index) + index = core.where(index < input.shape[dim], index, index - input.shape[dim]) + if not ON_ORANGE_PI: + return ops.gather_elements(input, dim, index) + return tf_gather(input, index, dim, batch_dims=dim) def gather_nd(input, indices): return ops.gather_nd(input, indices) @@ -195,10 +197,12 @@ def select(input, dim, index): # scatter has_scatter = hasattr(mindspore.mint, 'scatter') def scatter(input, dim, index, src): - if use_pyboost() and has_scatter: + if use_pyboost() and has_scatter and not ON_ORANGE_PI: return mindspore.mint.scatter(input, dim, index, src) if not isinstance(src, mindspore.Tensor): src = ops.full(index.shape, src, dtype=input.dtype) + if input.dtype == mindspore.bool_: + return ops.tensor_scatter_elements(input.int(), index, src.int(), dim).bool() return ops.tensor_scatter_elements(input, index, src, dim) def tf_scatter_nd_update(input, indices, updates): @@ -352,7 +356,7 @@ def _take_along_dim_helper(self, indices, dim): def take_along_dim(input, indices, dim=None, *, out=None): if dim: self_broadcasted, indices_broadcasted, dim = _take_along_dim_helper(input, indices, dim) - return self_broadcasted.gather(dim, indices_broadcasted) + return gather(self_broadcasted, dim, indices_broadcasted) return input.view(-1).gather(0, indices.view(-1)) # tensor_split @@ -427,7 +431,7 @@ def where(condition, *args, out=None): input = mindspore.tensor(input, dtype=mindspore.float32) other = finfo(input.dtype).min - if use_pyboost(): + if use_pyboost() and not ON_ORANGE_PI: output = mindspore.mint.where(condition, input, other) else: output = condition * input + (~condition) * other diff --git a/mindnlp/core/ops/comparison.py b/mindnlp/core/ops/comparison.py index 78cd343e0..18000f740 100644 --- a/mindnlp/core/ops/comparison.py +++ b/mindnlp/core/ops/comparison.py @@ -3,7 +3,7 @@ import numpy as np import mindspore from mindspore import ops -from ..configs import use_pyboost +from ..configs import use_pyboost, ON_ORANGE_PI from ._inner import call_ms_func @@ -64,7 +64,7 @@ def greater(input, other, *, out=None): # isclose has_isclose = hasattr(mindspore.mint, 'isclose') def isclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False): - if use_pyboost() and has_isclose: + if use_pyboost() and has_isclose and not ON_ORANGE_PI: return mindspore.mint.isclose(input, other, rtol, atol, equal_nan) return mindspore.tensor(np.isclose(input.numpy(), other.numpy(), rtol, atol, equal_nan)) @@ -174,7 +174,7 @@ def not_equal(input, other): # sort has_sort = hasattr(mindspore.mint, 'sort') def sort(input, *, dim=-1, descending=False, stable=False): - if use_pyboost() and has_sort: + if use_pyboost() and has_sort and not ON_ORANGE_PI: out = mindspore.mint.sort(input, dim=dim, descending=descending, stable=stable) else: out = ops.sort(input, dim, descending) diff --git a/mindnlp/core/ops/pointwise.py b/mindnlp/core/ops/pointwise.py index 8fe981939..98465b043 100644 --- a/mindnlp/core/ops/pointwise.py +++ b/mindnlp/core/ops/pointwise.py @@ -2,7 +2,7 @@ import numpy as np import mindspore from mindspore import ops -from ..configs import use_pyboost, ON_A1 +from ..configs import use_pyboost, ON_A1, ON_ORANGE_PI from ._inner import call_ms_func from mindnlp import core @@ -582,7 +582,7 @@ def igammac(input, other): def mul(input, other, *, out=None): - if use_pyboost() and has_mul: + if use_pyboost() and has_mul and not ON_ORANGE_PI: out = mindspore.mint.mul(input, other) else: if input.dtype == mindspore.bool_: diff --git a/mindnlp/utils/torch_proxy.py b/mindnlp/utils/torch_proxy.py index ae1fccae8..1882dc0b2 100644 --- a/mindnlp/utils/torch_proxy.py +++ b/mindnlp/utils/torch_proxy.py @@ -62,6 +62,8 @@ def exec_module(self, module): class ProxyModule(type(module)): def __getattr__(_, name): # 动态导入实际模块中的属性 + if DEVICE_TARGET == 'Ascend': + name = name.replace('cuda', 'npu') try: target_module = importlib.import_module(self.target_name) except ImportError as e: diff --git a/tests/run_test.py b/tests/run_test.py index 6891cf5da..6bd7ff7d4 100644 --- a/tests/run_test.py +++ b/tests/run_test.py @@ -1,6 +1,5 @@ import os import sys -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) import pytest import mindspore