use new safetensors loader (#1992)

Author: lvyufeng

.github/workflows/ci_pipeline.yaml

@@ -35,7 +35,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip==24.0
-        pip install -r requirements/requirements.txt
+        pip install -r requirements/pylint_requirements.txt
     - name: Install MindSpore
       shell: bash
       env:

llm/finetune/blenderbot/blenderbot_dolly-15k.py

@@ -8,12 +8,8 @@
 
 # 环境配置
 context.set_context(
-    mode=context.PYNATIVE_MODE,
-    device_target="Ascend",
     device_id=0,
-    enable_graph_kernel=False,
     max_call_depth=3000,
-    pynative_synchronize=True
 )
 ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.STAND_ALONE)
 ms.set_context(reserve_class_name_in_scope=False)

llm/finetune/blip2/blip2_finetune_with_Food500Cap.py

@@ -154,7 +154,7 @@ def eval(self):
         return vloss / step, testdatasetRES
 
 # 加载模型并设置可训练参数
-ms.set_context(device_target='Ascend', device_id=0, pynative_synchronize=True)
+ms.set_context(device_target='Ascend', device_id=0)
 processor = Blip2Processor.from_pretrained('Salesforce/blip2-opt-2.7b')
 model = Blip2ForConditionalGeneration.from_pretrained('Salesforce/blip2-opt-2.7b')
 model = freeze_blip2_backbone(model, freeze_vit=True)

llm/inference/baichuan/baichuan_13b.py

@@ -1,7 +1,6 @@
 import mindspore as ms
 from mindnlp.transformers import BaiChuanForCausalLM, BaiChuanTokenizer
 from mindnlp.transformers.generation.utils import GenerationConfig
-# ms.set_context(pynative_synchronize=True)
 
 tokenizer = BaiChuanTokenizer.from_pretrained("baichuan-inc/Baichuan-13B-Chat")
 model = BaiChuanForCausalLM.from_pretrained("baichuan-inc/Baichuan-13B-Chat",

llm/inference/llama2/simple_inference_with_static_cache.py

@@ -12,7 +12,7 @@
 
 model_id = 'shakechen/llama-2-7b-hf'
 tokenizer = LlamaTokenizer.from_pretrained(model_id, mirror='modelscope', pad_token="</s>", padding_side="right")
-model = LlamaForCausalLM.from_pretrained(model_id, mirror='modelscope', use_safetensors=False, ms_dtype=mindspore.float16)
+model = LlamaForCausalLM.from_pretrained(model_id, mirror='modelscope', use_safetensors=True, ms_dtype=mindspore.float16)
 inputs = tokenizer(prompts, return_tensors="ms", padding=True)
 
 def decode_one_tokens(model, cur_token, input_pos, cache_position, past_key_values):

mindnlp/core/nn/functional.py

@@ -13,7 +13,7 @@
 
 def gelu(input, approximate='none'):
     if use_pyboost():
-        return mindspore.mint.nn.functional.gelu(input, approximate)
+        return mindspore.mint.nn.functional.gelu(input, approximate=approximate)
     return ops.gelu(input, approximate)
 
 def relu(input):
@@ -111,8 +111,8 @@ def avg_pool1d(input_array, pool_size, stride, padding=0, ceil_mode=False, count
     output_array = ops.zeros((N, C, output_length))
 
     # Generate the starting indices of the pooling windows
-    indices = ops.arange(output_length) * stride
-    indices = indices[:, None] + ops.arange(pool_size)
+    indices = ops.range(0, output_length, 1) * stride
+    indices = indices[:, None] + ops.range(0, pool_size, 1)
 
     # Ensure indices are within bounds
     indices = ops.minimum(indices, input_array.shape[2] - 1)
@@ -1172,7 +1172,7 @@ def make_causal_mask(
     Returns:
       A `[batch..., 1, len, len]` shaped causal mask for 1d attention.
     """
-    idxs = ops.broadcast_to(ops.arange(x.shape[-1], dtype=mindspore.int32), x.shape)
+    idxs = ops.broadcast_to(ops.range(0, x.shape[-1], 1), x.shape)
     return make_attention_mask(
         idxs,
         idxs,

mindnlp/core/ops/creation.py

@@ -131,7 +131,11 @@ def empty(*size, dtype=None, device=None):
     if dtype is None:
         dtype = get_default_dtype()
     if has_empty:
-        out = mindspore._c_expression.pyboost_empty([size, dtype, device]).get_value()
+        out = mindspore._c_expression.pyboost_empty([size, dtype, device])
+        if not isinstance(out, mindspore.Tensor):
+            out = out.get_value()
+        else:
+            return out
     else:
         out = CTensor(dtype=dtype, shape=size)
     return mindspore.Tensor(out)

mindnlp/core/ops/reduction.py

@@ -48,24 +48,33 @@ def any(input, dim=None, keepdim=False):
 
 # max
 has_max = hasattr(mindspore.mint, 'max')
-def max(input, dim=None, keepdim=False):
+def max(*args, **kwargs):
     if use_pyboost() and has_max:
-        return mindspore.mint.max(input, dim, keepdim)
+        return mindspore.mint.max(*args, **kwargs)
+
+    input = kwargs.get('input', None) or args[0]
+    dim = kwargs.get('dim', None) or args[1]
+    keepdim = kwargs.get('keepdim', False) or args[2]
     out = ops.max(input, dim, keepdim)
     if dim is None:
         return out[0]
     return out
 
 # min
 has_min = hasattr(mindspore.mint, 'min')
-def min(input, dim=None, keepdim=False):
+def min(*args, **kwargs):
     if use_pyboost() and has_min:
-        return mindspore.mint.min(input, dim, keepdim)
+        return mindspore.mint.min(*args, **kwargs)
+
+    input = kwargs.get('input', None) or args[0]
+    dim = kwargs.get('dim', None) or args[1]
+    keepdim = kwargs.get('keepdim', False) or args[2]
     out = ops.min(input, dim, keepdim)
     if dim is None:
         return out[0]
     return out
 
+
 # dist
 
 # logsumexp

mindnlp/core/serialization.py

@@ -26,10 +26,8 @@
 import warnings
 import tempfile
 import operator
-import struct
 import mmap
 import json
-import math
 
 from contextlib import closing, contextmanager
 from enum import Enum
@@ -47,7 +45,6 @@
 
 import safetensors
 import safetensors.numpy
-from safetensors import deserialize
 
 from mindnlp.core import nn
 from mindnlp.core.nn import Parameter
@@ -65,6 +62,7 @@
 
 MAGIC_NUMBER = 0x1950A86A20F9469CFC6C
 PROTOCOL_VERSION = 1001
+MAX_HEADER_SIZE = 100 * 1000 * 1000
 
 
 @contextmanager
@@ -1433,50 +1431,148 @@ def _save(obj, zip_file, pickle_module, pickle_protocol):
 }
 
 
-def legacy_safe_load_file(filename):
-    """
-    This function safely loads a file containing state dictionary data and converts it into a dictionary of MindSpore Parameters.
+_DTYPE_SIZE = {
+    "BOOL": 1,
+    "U8": 1,
+    "I8": 1,
+    "F8_E5M2": 1,
+    "F8_E4M3": 1,
+    "I16": 2,
+    "U16": 2,
+    "I32": 4,
+    "U32": 4,
+    "I64": 8,
+    "U64": 8,
+    "F16": 2,
+    "BF16": 2,
+    "F32": 4,
+    "F64": 8,
+}
 
-    Args:
-        filename (str): The path to the file containing the state dictionary data to be loaded.
+class PySafeSlice:
+    def __init__(self, info, bufferfile, base_ptr, buffermmap):
+        self.info = info
+        self.bufferfile = bufferfile
+        self.buffermmap = buffermmap
+        self.base_ptr = base_ptr
+
+        self.start = [0 for _ in self.shape]
+        self.stop = list(self.shape)
+        self.step = [1 for _ in self.shape]
+
+    @property
+    def ndim(self):
+        return len(self.shape)
+
+    def get(self, *args, **kwargs):
+        nbytes = int(np.prod(self.shape)) * np.dtype(self.dtype).itemsize
+        offset = self.start_offset
+        tensor = np.frombuffer(self.buffermmap, dtype=self.dtype, offset=offset,
+                               count=nbytes // np.dtype(self.dtype).itemsize)
+        tensor = tensor.reshape(self.shape)
+        if not SUPPORT_BF16 and self.info["dtype"] == 'BF16':
+            tensor = tensor.astype(np.float16)
+        tensor = Tensor.from_numpy(tensor)
+        return tensor
+
+    @property
+    def start_offset(self):
+        return self.base_ptr + self.info["data_offsets"][0]
+
+    def get_shape(self):
+        return self.shape
+
+    @property
+    def shape(self):
+        return self.info["shape"]
+
+    @property
+    def dtype(self):
+        return _NP_TYPES[self.info["dtype"]]
+
+    @property
+    def nelements(self):
+        return np.prod(self.info["shape"])
+
+    @property
+    def bits(self):
+        return _DTYPE_SIZE[self.info["dtype"]]
+
+    @property
+    def nbytes(self):
+        return self.nelements * self.bits
+
+def getSize(fileobject):
+    fileobject.seek(0, 2)  # move the cursor to the end of the file
+    size = fileobject.tell()
+    fileobject.seek(0)  # move the cursor to the start of the file
+    return size
+
+
+def metadata_validate(metadata):
+    start = 0
+    for key, info in metadata.items():
+        s, e = info["data_offsets"]
+        if s != start or e < s:
+            raise ValueError(f"SafeTensorError::InvalidOffset({key})")
+        start = e
+        nelements = np.prod(info["shape"])
+        nbytes = nelements * _DTYPE_SIZE[info["dtype"]]
+        if (e - s) != nbytes:
+            raise ValueError("SafeTensorError::TensorInvalidInfo")
+    return start
+
+def read_metadata(buffer):
+    buffer_len = getSize(buffer)
+    if buffer_len < 8:
+        raise ValueError("SafeTensorError::HeaderTooSmall")
+
+    n = np.frombuffer(buffer.read(8), dtype=np.uint64).item()
+
+    if n > MAX_HEADER_SIZE:
+        raise ValueError("SafeTensorError::HeaderTooLarge")
+
+    stop = n + 8
+    if stop > buffer_len:
+        raise ValueError("SafeTensorError::InvalidHeaderLength")
+
+    tensors = json.loads(buffer.read(n), object_pairs_hook=OrderedDict)
+
+    metadata = tensors.pop("__metadata__", None)
+    buffer_end = metadata_validate(tensors)
+
+    if buffer_end + 8 + n != buffer_len:
+        raise ValueError("SafeTensorError::MetadataIncompleteBuffer")
+
+    return stop, tensors, metadata
+
+
+class fast_safe_open:
+    def __init__(self, filename, framework=None, device="cpu"):
+        self.filename = filename
+        self.framework = framework
+        self.file = open(self.filename, "rb")
+        self.file_mmap = mmap.mmap(self.file.fileno(), 0, access=mmap.ACCESS_COPY)
+        self.base, self.tensors_decs, self.__metadata__ = read_metadata(self.file)
+        self.tensors = OrderedDict()
+        for key, info in self.tensors_decs.items():
+            self.tensors[key] = PySafeSlice(info, self.file, self.base, self.file_mmap)
+            self.tensors[key].key = key
 
-    Returns:
-        dict: A dictionary where keys are parameter names and values are MindSpore Parameters.
+    def __enter__(self):
+        return self
 
-    Raises:
-        FileNotFoundError: If the specified file 'filename' does not exist.
-        ValueError: If the data in the file is not in the correct format to create MindSpore Parameters.
-    """
-    with open(filename, "rb") as f:
-        data = f.read()
+    def __exit__(self, *args):
+        self.file.close()
 
-    safeview = deserialize(data)
+    def metadata(self):
+        return self.__metadata__
 
-    result = {}
-    try:
-        for k, v in safeview:
-            dtype = _MS_TYPES[v["dtype"]]
-            if (not SUPPORT_BF16 and dtype != mindspore.bfloat16) or SUPPORT_BF16:
-                arr = Tensor.convert_bytes_to_tensor(
-                    bytes(v["data"]), tuple(v["shape"]), dtype
-                )
-                result[k] = Tensor(arr)
-            else:
-                raise TypeError(
-                    "Do not support bfloat16 on current device, use numpy as convert buffer to boost load."
-                )
-        return result
+    def keys(self):
+        return list(self.tensors.keys())
 
-    except Exception as e:
-        for k, v in safeview:
-            dtype = _NP_TYPES[v["dtype"]]
-            arr = np.frombuffer(v["data"], dtype=dtype).reshape(v["shape"])
-
-            if (not SUPPORT_BF16 and dtype != bfloat16) or SUPPORT_BF16:
-                result[k] = Tensor.from_numpy(arr)
-            else:
-                result[k] = Tensor.from_numpy(arr.astype(np.float16))
-        return result
+    def get_tensor(self, name):
+        return self.tensors[name].get()
 
 
 def safe_load_file(filename):
@@ -1494,39 +1590,10 @@ def safe_load_file(filename):
         ValueError: If the data in the file is not in the correct format to create MindSpore Parameters.
     """
 
-    def convert(info: dict[str, Any]):
-        numpy_dtype = _NP_TYPES[info["dtype"]]
-        ms_dtype = _MS_TYPES[info["dtype"]]
-        shape: list[int] = info["shape"]
-        begin, end = info["data_offsets"]
-        assert 0 <= begin <= end <= len(byte_buf)
-        assert end - begin == math.prod(shape) * np.dtype(numpy_dtype).itemsize
-        buf = byte_buf[begin:end]
-
-        array = np.frombuffer(buf, dtype=numpy_dtype).reshape(shape)
-
-        if array.dtype == bfloat16 and not SUPPORT_BF16:
-            logger.warning_once(
-                "MindSpore do not support bfloat16 dtype, we will automaticlly convert to float16"
-            )
-            array = array.astype(np.float16)
-        array = array.astype(array.dtype)
-        out = Tensor.from_numpy(array)
-        return out
-
-    with open(filename, "rb") as fp:
-        (header_size,) = struct.unpack("<Q", fp.read(8))
-        header: dict[str, dict[str, Any]] = json.loads(fp.read(header_size))
-        # Use mmap for the actual data to avoid race conditions with the file offset.
-        mapped = memoryview(mmap.mmap(fp.fileno(), 0, access=mmap.ACCESS_READ))
-        byte_buf = mapped[8 + header_size :]
-
-        result = {
-            name: convert(info)
-            for (name, info) in header.items()
-            if name != "__metadata__"
-        }
-
+    result = {}
+    with fast_safe_open(filename, framework="np") as f:
+        for k in f.keys():
+            result[k] = f.get_tensor(k)
     return result
 
 

mindnlp/transformers/models/llama/modeling_llama.py

@@ -98,10 +98,10 @@ def _prepare_4d_causal_attention_mask_with_cache_position(
             #     padding_mask, min_dtype
             # )
             if mask_length >= causal_mask.shape[-1]:
-                causal_mask = causal_mask.masked_fill(padding_mask, min_dtype)
+                causal_mask = ops.masked_fill(causal_mask, padding_mask, min_dtype)
             else:
                 causal_mask = ops.cat(
-                                [ops.narrow(causal_mask, -1, 0, mask_length).masked_fill(padding_mask, min_dtype),
+                                [ops.masked_fill(ops.narrow(causal_mask, -1, 0, mask_length), padding_mask, min_dtype),
                                 ops.narrow(causal_mask, -1, mask_length, causal_mask.shape[-1] - mask_length)],
                                 dim=-1
                             )