mindspore 2.4.10__cp311-cp311-win_amd64.whl → 2.6.0rc1__cp311-cp311-win_amd64.whl

mindspore/experimental/es/embedding_service.py

@@ -125,18 +125,20 @@ class EmbeddingServiceOut:
 
 class EmbeddingService:
     r"""
-    Currently, ES(EmbeddingService) feature can only create one object which can support model training and inference
+    ES(EmbeddingService) feature can support model training and inference
     for PS embedding and data_parallel embedding, and provide unified embedding management, storage,
     and computing capabilities for training and inference.
     PS embedding refer to tables that vocab_size more than 100,000, and recommended to store them on the
     Parameter Server (PS). Data_parallel embedding refer to tables that vocab_size less than 100,000, and recommended
     to store them on device.
 
+    Currently, ES feature can only create one instance of EmbeddingService object.
+
     .. warning::
         This is an experimental EmbeddingService API that is subject to change.
 
     .. note::
-        This API needs to call 'mindspore.communication.init()' before,
+        This API needs to call :func:`mindspore.communication.init` before,
         and it can take effect after the dynamic networking is completed.
 
     Raises:
@@ -241,24 +243,26 @@ class EmbeddingService:
             name (str): The embedding table name.
             init_vocabulary_size (int): The size of embedding table.
             embedding_dim (int): The embedding dim of data in embedding table.
-            max_feature_count (int): The count of keys when look up for PS.
-            initializer (Initializer): The initialization strategy for the PS embedding, default is ``Uniform``.
-            embedding_type (str): The embedding type, configurable parameters ["PS", "data_parallel"],
+            max_feature_count (int, optional): The count of keys when look up for PS. Default: ``None``.
+            initializer (Initializer, optional): The initialization strategy for the PS embedding,
+                default is ``Uniform(scale=0.01)``.
+            embedding_type (str, optional): The embedding type, configurable parameters ["PS", "data_parallel"],
                 ``"PS"`` means initializing PS embedding, ``"data_parallel"`` means initializing data_parallel
                 embedding, and default is ``"PS"``.
-            ev_option (EmbeddingVariableOption): Properties of the PS embedding,
+            ev_option (EmbeddingVariableOption, optional): Properties of the PS embedding,
                 is a EmbeddingVariableOption obj which returned by embedding_variable_option function.
                 Default is ``None``.
-            multihot_lens (int): The param only use when allow_merge is enabled, and not support now.
+            multihot_lens (int, optional): The param only use when `allow_merge` is enabled, and not support now.
                 Default is ``None``.
-            optimizer (str): The type of optimizer in the train mode for PS embedding,
+            optimizer (str, optional): The type of optimizer in the train mode for PS embedding,
                 cannot be shared among each PS embedding, and currently only ``"Adam"``, ``"Ftrl"``, ``"SGD"`` and
                 ``"RMSProp"`` are supported, and default is ``None``.
-            allow_merge (bool): Whether to enable merge data_parallel embeddings, currently only be False,
+            allow_merge (bool, optional): Whether to enable merge data_parallel embeddings, currently only be False,
                 and default is ``False``.
-            optimizer_param (float): The "initialize accumulator value" param of optimizer which configured by user,
+            optimizer_param (float, optional): The "initialize accumulator value" param
+                of optimizer which configured by user,
                 representing the init value of moment accumulator, and default is ``None``.
-            mode (str): Run mode, configurable parameters ["train", "predict", "export"],
+            mode (str, optional): Run mode, configurable parameters ["train", "predict", "export"],
                 ``"train"`` means train mode, ``"predict"`` means predict mode, ``"export"`` mean export mode,
                 and default is ``"train"``.
 
@@ -345,8 +349,9 @@ class EmbeddingService:
 
         Args:
             padding_key (int): The value for padding key, must be a genuine and legal hash key.
-            mask (bool): Whether to update padding key. If set to false, it will not be updated. Default is ``True``.
-            mask_zero (bool): Whether to update padding key when key is 0. Default is ``False``.
+            mask (bool, optional): Whether to update padding key. If set to false, it will not be updated.
+                Default is ``True``.
+            mask_zero (bool, optional): Whether to update padding key when key is 0. Default is ``False``.
 
         Returns:
             PaddingParamsOption object.
@@ -368,7 +373,7 @@ class EmbeddingService:
 
         Args:
             completion_key (int): The value for completion key.
-            mask (bool): Whether to update completion key. If set to false, it will not be updated,
+            mask (bool, optional): Whether to update completion key. If set to false, it will not be updated,
                 and default is ``True``.
 
         Returns:
@@ -396,10 +401,11 @@ class EmbeddingService:
 
         Args:
             filter_freq (int): The frequency threshold value for feature admission.
-            default_key (int): The key that number of occurrences does not reach the threshold,
-                return value of default key as the corresponding value when look up embedding, and default is ``None``.
-            default_value (int/float): The key that number of occurrences does not reach the threshold,
-                return default value which length value is embedding dim, and default is ``None``.
+            default_key (int, optional): The key that number of occurrences does not reach the threshold,
+                return value of `default_key` as the corresponding value when look up embedding,
+                and default is ``None``.
+            default_value (Union[int, float], optional): The key that number of occurrences does not
+                reach the threshold, return default value which length value is embedding dim, and default is ``None``.
 
         Returns:
             CounterFilter object.
@@ -460,16 +466,17 @@ class EmbeddingService:
         Set variable option for PS embedding.
 
         Args:
-            filter_option (CounterFilter): The option of counter filter. Default is ``None``.
-            padding_option (PaddingParamsOption): The option of padding key. Default is ``None``.
-            evict_option (EvictOption): The option evict. Default is ``None``.
-            completion_option (CompletionKeyOption): The option of completion key. Default is ``None``.
-            storage_option (None): Reserved option, currently not supported. Default is ``None``.
-            feature_freezing_option (None): Reserved option, currently not supported. Default is ``None``.
-            communication_option (None): Reserved option, currently not supported. Default is ``None``.
+            filter_option (CounterFilter, optional): The option of counter filter. Default is ``None``.
+            padding_option (PaddingParamsOption, optional): The option of padding key. Default is ``None``.
+            evict_option (EvictOption, optional): The option evict. Default is ``None``.
+            completion_option (CompletionKeyOption, optional): The option of completion key. Default is ``None``.
+            storage_option (None, optional): Reserved option, currently not supported. Default is ``None``.
+            feature_freezing_option (None, optional): Reserved option, currently not supported. Default is ``None``.
+            communication_option (None, optional): Reserved option, currently not supported. Default is ``None``.
 
         Returns:
-            EmbeddingVariableOption object, used as the ev_option parameter for embedding_init.
+            EmbeddingVariableOption object, used as the ev_option parameter for
+            :func:`mindspore.experimental.es.EmbeddingService.embedding_init` .
 
         Raises:
             TypeError: If value of "filter_option" is not None and the type of "filter_option" is not CounterFilter.
@@ -501,7 +508,8 @@ class EmbeddingService:
 
         .. note::
             This function can only be executed by rank 0.
-            Need to call embedding_variable_option to set evict_option for each PS embedding before export.
+            Need to call :func:`mindspore.experimental.es.EmbeddingService.embedding_variable_option`
+            to set evict_option for each PS embedding before export.
 
         Args:
             file_path (str): The path to export embedding ckpt, and the last character cannot be ``"/"``.

mindspore/experimental/llm_boost/__init__.py

@@ -16,6 +16,7 @@
 from __future__ import absolute_import
 
 from mindspore.experimental.llm_boost.atb import LlamaBoost, QwenBoost
+from mindspore.experimental.llm_boost.ascend_native import *
 from mindspore.experimental.llm_boost.register import LlmBoostRegister
 
 __all__ = ["LlmBoostRegister"]

mindspore/experimental/llm_boost/ascend_native/__init__.py

@@ -0,0 +1,22 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""
+Provide llm boost for inference, such as LlamaBoost.
+"""
+from __future__ import absolute_import
+
+from mindspore.experimental.llm_boost.ascend_native.llama_boost_ascend_native import LlamaBoostAscendNative
+
+__all__ = ['LlamaBoostAscendNative']

mindspore/experimental/llm_boost/ascend_native/llama_boost_ascend_native.py

@@ -0,0 +1,211 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""AscendNative Llama Boost APIs."""
+
+import os
+import numpy as np
+from mindspore.common import Tensor, dtype
+from mindspore.experimental.llm_boost.ascend_native.llm_boost import LLMBoost
+from mindspore.experimental.llm_boost.register import LlmBoostRegister, LlmBoostType
+
+
+def RoundUp(val: int, align: int) -> int:
+    if align == 0:
+        return 0
+    return -(val // -align) * align
+
+
+def ConvertTensor(nd_mat: np.ndarray, transpose: bool = True, nd2nz: bool = True) -> np.ndarray:
+    """ Transforms tensor format from Nd to Nz """
+    if transpose:
+        nd_mat = np.transpose(nd_mat)
+    if not nd2nz:
+        return nd_mat
+    block_size = (16, 16)
+    r = RoundUp(nd_mat.shape[0], block_size[0])
+    c = RoundUp(nd_mat.shape[1], block_size[1])
+    r_pad = r - nd_mat.shape[0]
+    c_pad = c - nd_mat.shape[1]
+    nd_mat = np.pad(nd_mat, ((0, r_pad), (0, c_pad)))
+    nz_mat = np.transpose(np.reshape(
+        nd_mat, (r, c // block_size[1], block_size[1])), (1, 0, 2))
+    nz_mat = nz_mat.reshape(r, c)
+    return nz_mat
+
+
+@LlmBoostRegister.register(LlmBoostType.ASCEND_NATIVE, "Llama")
+class LlamaBoostAscendNative(LLMBoost):
+    r"""
+    Implements an Llama model in a single kernel.
+    it forwards the python functions to the C++ binded object
+    """
+    def _get_from_dict(self, dictionary, name):
+        """ internal function to get a specific tensor from the dictionary """
+        all_relevant_layers = [value for key, value in dictionary.items() if name in key]
+        if all_relevant_layers:
+            return all_relevant_layers[0].asnumpy()
+        return None
+
+    def _get_quant_triplet_from_dict(self, dictionary, name):
+        """ internal function to get a weight triple tensor from the dictionary """
+        weights = self._get_from_dict(dictionary, name + "._handler.weight")
+        scale = self._get_from_dict(dictionary, name + "._weight_quantizer.scale")
+        offset = self._get_from_dict(dictionary, name + "._weight_quantizer.zp_neg")
+        return weights, scale, offset
+
+    def _prepare_single_layer(self, ckpt, config, id):
+        """ prepares the dictionary of weights of a single layer """
+        prefix = 'model.layers.' + str(id)
+        is_last = (id == config.num_layers-1)
+        layer = 'layers.' + str(id) + '.'
+        l_dict = {key: value for key, value in ckpt.items() if layer in key}
+        if config.n_kv_heads is None:
+            config.n_kv_heads = config.num_heads
+        start = 0
+        end = config.hidden_size
+        kv_start = 0
+        kv_end = int(config.hidden_size*config.n_kv_heads/config.num_heads)
+        ffn_hid = [value for key, value in l_dict.items() if "w3" in key][0].shape[0]
+        ffn_start = 0
+        ffn_end = ffn_hid
+        rank_size = int(os.getenv('RANK_SIZE', '1'))
+        #Emir if (config.parallel_mode != 2): # 2 - AUTO_PARALLEL
+        hid_size = end
+        kv_hid_size = kv_end
+        embed_size = config.vocab_size
+        rank_id = int(os.getenv('RANK_ID', '0'))
+        if (hid_size % rank_size == 0) and (ffn_hid % rank_size == 0) and (embed_size % rank_size == 0):
+            start = int(rank_id * hid_size / rank_size)
+            end = int((rank_id + 1) * hid_size / rank_size)
+            kv_start = int(rank_id * kv_hid_size / rank_size)
+            kv_end = int((rank_id + 1) * kv_hid_size / rank_size)
+            ffn_start = int(rank_id * ffn_hid / rank_size)
+            ffn_end = int((rank_id + 1) * ffn_hid / rank_size)
+        else:
+            raise RuntimeError("hidden size and ffn hidden size must be divided by rank size without remainder.  \
+                                hidden_size: ", hid_size, " ffn_hidden_size: ", ffn_hid, " rank_size: ", rank_size)
+        quant = (self._get_from_dict(l_dict, "_weight_quantizer") is not None)
+        unite_qkv = (config.num_heads == config.n_kv_heads)
+        self.dictionary[prefix + ".attention_norm.weight"] = \
+            Tensor(self._get_from_dict(l_dict, "attention_norm"), dtype=dtype.float16)
+        self.dictionary[prefix + ".ffn_norm.weight"] = \
+            Tensor(self._get_from_dict(l_dict, "ffn_norm"), dtype=dtype.float16)
+        if is_last:
+            self.dictionary['lm_head.weight'] = Tensor(ConvertTensor(ckpt['lm_head.weight'].asnumpy()[:, start:end]))
+
+        if not quant:
+            self._pack_attn_weights(l_dict, prefix, start, end, kv_start, kv_end, unite_qkv)
+            self._pack_ffn_weights(l_dict, prefix, ffn_start, ffn_end)
+        else:
+            self._pack_attn_quant_weights(l_dict, prefix, start, end, kv_start, kv_end, unite_qkv)
+            self._pack_ffn_quant_weights(l_dict, prefix, ffn_start, ffn_end)
+
+    def _pack_attn_weights(self, l_dict, prefix, start, end, kv_start, kv_end, unite_qkv):
+        """ prepares the dictionary of weights of an attention block """
+        wq = self._get_from_dict(l_dict, "wq")[start:end, :]
+        wk = self._get_from_dict(l_dict, "wk")[kv_start:kv_end, :]
+        wv = self._get_from_dict(l_dict, "wv")[kv_start:kv_end, :]
+        self.dictionary[prefix + ".attention.wo.weight"] = \
+            Tensor(ConvertTensor(self._get_from_dict(l_dict, "wo")[:, start:end]))
+        if unite_qkv:
+            self.dictionary[prefix + ".attention.wqkv.weight"] = Tensor(ConvertTensor(np.concatenate((wq, wk, wv))))
+        else:
+            self.dictionary[prefix + ".attention.wq.weight"] = Tensor(ConvertTensor(wq))
+            self.dictionary[prefix + ".attention.wkv.weight"] = Tensor(ConvertTensor(np.concatenate((wk, wv))))
+
+    def _pack_ffn_weights(self, l_dict, prefix, ffn_start, ffn_end):
+        """ prepares the dictionary of weights of an ffn block """
+        self.dictionary[prefix + ".feed_forward.w2.weight"] = \
+            Tensor(ConvertTensor(self._get_from_dict(l_dict, "w2")[:, ffn_start:ffn_end]))
+        w1 = self._get_from_dict(l_dict, "w1")[ffn_start:ffn_end, :]
+        w3 = self._get_from_dict(l_dict, "w3")[ffn_start:ffn_end, :]
+        self.dictionary[prefix + ".feed_forward.w13.weight"] = Tensor(ConvertTensor(np.concatenate((w1, w3))))
+
+    def _pack_attn_quant_weights(self, l_dict, prefix, start, end, kv_start, kv_end, unite_qkv):
+        """ prepares the dictionary of weights of a quantized attention block """
+        wq, wq_scale, wq_offset = self._get_quant_triplet_from_dict(l_dict, "wq")
+        wk, wk_scale, wk_offset = self._get_quant_triplet_from_dict(l_dict, "wk")
+        wv, wv_scale, wv_offset = self._get_quant_triplet_from_dict(l_dict, "wv")
+        wo, wo_scale, wo_offset = self._get_quant_triplet_from_dict(l_dict, "wo")
+        self.dictionary[prefix + ".attention.wo.weight"] = Tensor(ConvertTensor(wo[:, start:end], nd2nz=False))
+        self.dictionary[prefix + ".attention.wo.weight.scale"] = Tensor(wo_scale[start:end])
+        self.dictionary[prefix + ".attention.wo.weight.offset"] = Tensor(wo_offset[start:end])
+
+        if unite_qkv:
+            self.dictionary[prefix + ".attention.wqkv.weight"] = \
+             Tensor(ConvertTensor(np.concatenate((wq[start:end, :], wk[kv_start:kv_end, :], wv[kv_start:kv_end, :])),
+                                  nd2nz=False))
+            self.dictionary[prefix + ".attention.wqkv.weight.scale"] = \
+                Tensor(np.concatenate((wq_scale[start:end], wk_scale[kv_start:kv_end], wv_scale[kv_start:kv_end])))
+            self.dictionary[prefix + ".attention.wqkv.weight.offset"] = \
+                Tensor(np.concatenate((wq_offset[start:end], wk_offset[kv_start:kv_end], wv_offset[kv_start:kv_end])))
+        else:
+            self.dictionary[prefix + ".attention.wq.weight"] = Tensor(ConvertTensor(wq[start:end, :], nd2nz=False))
+            self.dictionary[prefix + ".attention.wq.weight.scale"] = Tensor(wq_scale[start:end])
+            self.dictionary[prefix + ".attention.wq.weight.offset"] = Tensor(wq_offset[start:end])
+            self.dictionary[prefix + ".attention.wkv.weight"] = \
+                Tensor(ConvertTensor(np.concatenate((wk[kv_start:kv_end, :], wv[kv_start:kv_end, :])), nd2nz=False))
+            self.dictionary[prefix + ".attention.wkv.weight.scale"] = \
+                Tensor(np.concatenate((wk_scale[kv_start:kv_end], wv_scale[kv_start:kv_end])))
+            self.dictionary[prefix + ".attention.wkv.weight.offset"] = \
+                Tensor(np.concatenate((wk_offset[kv_start:kv_end], wv_offset[kv_start:kv_end])))
+
+    def _pack_ffn_quant_weights(self, l_dict, prefix, ffn_start, ffn_end):
+        """ prepares the dictionary of weights of a quantized ffn block """
+        w1, w1_scale, w1_offset = self._get_quant_triplet_from_dict(l_dict, "w1")
+        w2, w2_scale, w2_offset = self._get_quant_triplet_from_dict(l_dict, "w2")
+        w3, w3_scale, w3_offset = self._get_quant_triplet_from_dict(l_dict, "w3")
+        self.dictionary[prefix + ".feed_forward.w2.weight"] = Tensor(ConvertTensor(w2[:, ffn_start:ffn_end],
+                                                                                   nd2nz=False))
+        self.dictionary[prefix + ".feed_forward.w2.weight.scale"] = Tensor(w2_scale[ffn_start:ffn_end])
+        self.dictionary[prefix + ".feed_forward.w2.weight.offset"] = Tensor(w2_offset[ffn_start:ffn_end])
+
+        self.dictionary[prefix + ".feed_forward.w13.weight"] = \
+                Tensor(ConvertTensor(np.concatenate((w1[ffn_start:ffn_end, :], w3[ffn_start:ffn_end, :])), nd2nz=False))
+        self.dictionary[prefix + ".feed_forward.w13.weight.scale"] = \
+                Tensor(np.concatenate((w1_scale[ffn_start:ffn_end], w3_scale[ffn_start:ffn_end])))
+        self.dictionary[prefix + ".feed_forward.w13.weight.offset"] = \
+            Tensor(np.concatenate((w1_offset[ffn_start:ffn_end], w3_offset[ffn_start:ffn_end])))
+
+    def _prepare_cos_sin_arrays(self, config, theta=10000):
+        """ prepares the cosine and sine arrays """
+        head_dim = config.hidden_size // config.num_heads
+        max_position_embedding = \
+            config.max_position_embedding if config.max_position_embedding is not None else config.seq_length
+        freqs_base = np.arange(0, head_dim, 2)[: (head_dim // 2)].astype(np.float32)
+        freqs = 1.0 / (theta ** (freqs_base / head_dim))
+        t = np.arange(0, max_position_embedding, 1).astype(np.float32)
+        freqs = np.outer(t, freqs)
+        emb = np.concatenate((freqs, freqs), axis=-1)
+        freqs_cos = Tensor(np.cos(emb), dtype=dtype.float16)
+        sin = np.sin(emb)
+
+        sin[:, :int(emb.shape[1]/2)] = -sin[:, :int(emb.shape[1]/2)]
+        self.dictionary['model.cos.weight'] = freqs_cos
+        freqs_sin = Tensor(sin, dtype=dtype.float16)
+        self.dictionary['model.sin.weight'] = freqs_sin
+
+    def set_weights(self, ckpt_dict):
+        """ load the checkpoint """
+        self.dictionary = {}
+        self.dictionary['model.tok_embeddings.embedding_weight'] = \
+            Tensor(ckpt_dict['model.tok_embeddings.embedding_weight'].asnumpy())
+        self.dictionary['model.norm_out.weight'] = \
+            Tensor(ckpt_dict['model.norm_out.weight'].asnumpy(), dtype=dtype.float16)
+        self._prepare_cos_sin_arrays(self.config)
+        for layer_id in range(self.config.num_layers):
+            self._prepare_single_layer(ckpt_dict, self.config, layer_id)
+
+        self.binder.set_weights_map(self.dictionary)

mindspore/experimental/llm_boost/ascend_native/llm_boost.py

@@ -0,0 +1,52 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""LLMBoost APIs."""
+
+from mindspore.common import Tensor
+
+class LLMBoost():
+    r"""
+    Implements an LLM in a single kernel.
+    it forwards the python function to the C++ binded object
+    """
+    def __init__(self, config):
+        r"""
+        initialize the parameters of the llm binder.
+        config is simply the config object of the model
+        """
+        from mindspore._c_expression import LlmBoostBinder
+        self.config = config
+        self.binder = LlmBoostBinder("AscendNative", config.model_type)
+        self.binder.init_model(config.to_dict())
+
+    def init(self):
+        """
+        Initialize the object
+        returns True if object needs input manipulation by mindformers
+        """
+        return False
+
+    def set_kvcache(self, k_caches=None, v_caches=None):
+        return
+
+    def forward(self, input_ids, batch_valid_length, position_ids=None):
+        ret = self.binder.forward([input_ids, batch_valid_length], "nothing really")
+        return Tensor(ret[0])
+
+    def set_weights(self, ckpt_dict):
+        self.binder.set_weights_map(ckpt_dict)
+
+    def add_flags(self, is_first_iteration=False):
+        self.binder.add_flags(is_first_iteration=is_first_iteration)

mindspore/experimental/llm_boost/atb/boost_base.py

@@ -112,8 +112,7 @@ class AtbBoostBase:
 
     def _convert_qkv_concat_weight(self, param_dict):
         """convert qkv concat weight"""
-        assume_num_layers = 500
-        for i in range(assume_num_layers):
+        for i in range(self.num_layers):
             # qkv weight concat
             wq_weight_name = f"model.layers.{i}.attention.wq.weight"
             wk_weight_name = f"model.layers.{i}.attention.wk.weight"
@@ -151,7 +150,7 @@ class AtbBoostBase:
             logger.info(f"transform: {qkv_concat_weight_name}")
             logger.info(f"transform: {gate_hidden_concat_weight_name}")
 
-        for i in range(assume_num_layers):
+        for i in range(self.num_layers):
             # qkv bias concat
             wq_bias_name = f"model.layers.{i}.attention.wq.bias"
             wk_bias_name = f"model.layers.{i}.attention.wk.bias"

mindspore/experimental/llm_boost/atb/llama_boost.py

@@ -43,7 +43,11 @@ class LlamaBoost(AtbBoostBase):
         )
 
     def init(self):
-        """set param"""
+        """
+        Initialize the object
+        returns True if object needs input manipulation by mindformers
+        """
+
         coder_param = {
             "normEps": self.config.rms_norm_eps,
             "normType": NormType.RMS_NORM,
@@ -93,6 +97,7 @@ class LlamaBoost(AtbBoostBase):
         }
         self.atb_encoder_operation.init(json.dumps({**encoder_param}))
         self.atb_decoder_operation.init(json.dumps({**decoder_param}))
+        return True
 
     def _prepare_inputs(
             self,

mindspore/experimental/llm_boost/register.py

@@ -23,6 +23,7 @@ class LlmBoostType:
         pass
 
     BUILDIN = 'BuildIn'
+    ASCEND_NATIVE = 'LLMBoost'
 
 
 class LlmBoostRegister:

mindspore/experimental/map_parameter.py

@@ -23,7 +23,7 @@ from copy import copy
 import numbers
 import mindspore as ms
 from mindspore.common.parameter import Parameter, _get_unique_parameter_key
-from mindspore._c_expression import Tensor as Tensor_
+from mindspore._c_expression import TensorPy as Tensor_
 from mindspore._c_expression import MapTensor_
 from mindspore.ops.operations import _map_tensor_ops
 
@@ -78,12 +78,12 @@ class MapParameter(Parameter):
         if value_dtype is not None:
             if isinstance(value_shape, numbers.Number):
                 value_shape = (value_shape,)
-            data = Tensor_(value_dtype, value_shape)
+            data = Tensor_(dtype=value_dtype, shape=value_shape)
         elif value_tensor is not None:
-            data = Tensor_(value_tensor.dtype, value_tensor.shape)
+            data = Tensor_(dtype=value_tensor.dtype, shape=value_tensor.shape)
         else:
             # default
-            data = Tensor_(ms.float32, (1,))
+            data = Tensor_(dtype=ms.float32, shape=(1,))
         obj = Tensor_.__new__(cls)
         Tensor_.__init__(obj, data)
         # Compatible attributes with Parameter.

mindspore/experimental/optim/adadelta.py

@@ -37,14 +37,14 @@ class Adadelta(Optimizer):
     Implements Adadelta algorithm.
 
     .. math::
-       \begin{aligned}
-            &\rule{150mm}{0.4pt}                                                                 \\
+        \begin{aligned}
+            &\rule{180mm}{0.4pt}                                                                 \\
             &\textbf{input}      : \gamma \text{ (lr)}, \: \theta_0 \text{ (params)},
                 \: f(\theta) \text{ (objective)}, \: \rho \text{ (decay)},
                 \: \lambda \text{ (weight decay)}                                                \\
             &\textbf{initialize} :  v_0  \leftarrow 0 \: \text{ (square avg)},
                 \: u_0 \leftarrow 0 \: \text{ (accumulate variables)}                     \\[-1.ex]
-            &\rule{110mm}{0.4pt}                                                                 \\
+            &\rule{180mm}{0.4pt}                                                                 \\
             &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}                         \\
             &\hspace{5mm}g_t           \leftarrow   \nabla_{\theta} f_t (\theta_{t-1})           \\
             &\hspace{5mm}if \: \lambda \neq 0                                                    \\
@@ -55,10 +55,10 @@ class Adadelta(Optimizer):
             &\hspace{5mm} u_t  \leftarrow   u_{t-1}  \rho +
                  \Delta x^2_t  (1 - \rho)                                                        \\
             &\hspace{5mm}\theta_t      \leftarrow   \theta_{t-1} - \gamma  \Delta x_t            \\
-            &\rule{110mm}{0.4pt}                                                          \\[-1.ex]
+            &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
             &\bf{return} \:  \theta_t                                                     \\[-1.ex]
-            &\rule{110mm}{0.4pt}                                                          \\[-1.ex]
-       \end{aligned}
+            &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
+        \end{aligned}
 
     .. warning::
         This is an experimental optimizer API that is subject to change.

mindspore/experimental/optim/adagrad.py

@@ -38,12 +38,12 @@ class Adagrad(Optimizer):
 
     .. math::
        \begin{aligned}
-            &\rule{110mm}{0.4pt}                                                                 \\
+            &\rule{160mm}{0.4pt}                                                                 \\
             &\textbf{input}      : \gamma \text{ (lr)}, \: \theta_0 \text{ (params)}, \: f(\theta)
                 \text{ (objective)}, \: \lambda \text{ (weight decay)},                          \\
             &\hspace{12mm}    \tau \text{ (initial accumulator value)}, \: \eta\text{ (lr decay)}\\
             &\textbf{initialize} :  state\_sum_0 \leftarrow 0                             \\[-1.ex]
-            &\rule{110mm}{0.4pt}                                                                 \\
+            &\rule{160mm}{0.4pt}                                                                 \\
             &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}                         \\
             &\hspace{5mm}g_t           \leftarrow   \nabla_{\theta} f_t (\theta_{t-1})           \\
             &\hspace{5mm} \tilde{\gamma}    \leftarrow \gamma / (1 +(t-1) \eta)                  \\
@@ -52,9 +52,9 @@ class Adagrad(Optimizer):
             &\hspace{5mm}state\_sum_t  \leftarrow  state\_sum_{t-1} + g^2_t                      \\
             &\hspace{5mm}\theta_t \leftarrow
                 \theta_{t-1}- \tilde{\gamma} \frac{g_t}{\sqrt{state\_sum_t}+\epsilon}            \\
-            &\rule{110mm}{0.4pt}                                                          \\[-1.ex]
+            &\rule{160mm}{0.4pt}                                                          \\[-1.ex]
             &\bf{return} \:  \theta_t                                                     \\[-1.ex]
-            &\rule{110mm}{0.4pt}                                                          \\[-1.ex]
+            &\rule{160mm}{0.4pt}                                                          \\[-1.ex]
        \end{aligned}
 
     .. warning::

mindspore/experimental/optim/adam.py

@@ -49,12 +49,14 @@ class Adam(Optimizer):
 
     .. math::
         \begin{aligned}
+            &\rule{180mm}{0.4pt}                                                                 \\
             &\textbf{input}      : \gamma \text{ (lr)}, \beta_1, \beta_2
                 \text{ (betas)},\theta_0 \text{ (params)},f(\theta) \text{ (objective)}          \\
             &\hspace{13mm}      \lambda \text{ (weight decay)},  \: \textit{amsgrad},
                 \:\textit{maximize}                                                              \\
             &\textbf{initialize} :  m_0 \leftarrow 0 \text{ ( first moment)},
                 v_0\leftarrow 0 \text{ (second moment)},\: \widehat{v_0}^{max}\leftarrow 0\\[-1.ex]
+            &\rule{180mm}{0.4pt}                                                                 \\
             &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}                         \\
             &\hspace{5mm}\textbf{if} \: \textit{maximize}:                                       \\
             &\hspace{10mm}g_t           \leftarrow   -\nabla_{\theta} f_t (\theta_{t-1})         \\
@@ -74,10 +76,15 @@ class Adam(Optimizer):
             &\hspace{5mm}\textbf{else}                                                           \\
             &\hspace{10mm}\theta_t \leftarrow \theta_{t-1} - \gamma \widehat{m_t}/
                 \big(\sqrt{\widehat{v_t}} + \epsilon \big)                                       \\
+            &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
             &\bf{return} \:  \theta_t                                                     \\[-1.ex]
+            &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
        \end{aligned}
 
     .. warning::
+        The implementation formula of this optimizer interface is not completely consistent with that in the paper.
+        If you want to use an interface that is completely consistent, it is recommended to use
+        :class:`mindspore.mint.optim.Adam`, which currently only supports Ascend.
         This is an experimental optimizer API that is subject to change.
         This module must be used with lr scheduler module in `LRScheduler Class
         <https://www.mindspore.cn/docs/en/master/api_python/mindspore.nn.html#learningrateschedule-class>`_ .

mindspore/experimental/optim/adamax.py

@@ -43,14 +43,14 @@ class Adamax(Optimizer):
 
     .. math::
        \begin{aligned}
-            &\rule{110mm}{0.4pt}                                                                 \\
+            &\rule{180mm}{0.4pt}                                                                 \\
             &\textbf{input}      : \gamma \text{ (lr)}, \beta_1, \beta_2
                 \text{ (betas)},\theta_0 \text{ (params)},f(\theta) \text{ (objective)},
                 \: \lambda \text{ (weight decay)},                                                \\
             &\hspace{13mm}    \epsilon \text{ (epsilon)}                                          \\
             &\textbf{initialize} :  m_0 \leftarrow 0 \text{ ( first moment)},
                 u_0 \leftarrow 0 \text{ ( infinity norm)}                                 \\[-1.ex]
-            &\rule{110mm}{0.4pt}                                                                 \\
+            &\rule{180mm}{0.4pt}                                                                 \\
             &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}                         \\
             &\hspace{5mm}g_t           \leftarrow   \nabla_{\theta} f_t (\theta_{t-1})           \\
             &\hspace{5mm}if \: \lambda \neq 0                                                    \\
@@ -58,9 +58,9 @@ class Adamax(Optimizer):
             &\hspace{5mm}m_t      \leftarrow   \beta_1 m_{t-1} + (1 - \beta_1) g_t               \\
             &\hspace{5mm}u_t      \leftarrow   \mathrm{max}(\beta_2 u_{t-1}, |g_{t}|+\epsilon)   \\
             &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \frac{\gamma m_t}{(1-\beta^t_1) u_t} \\
-            &\rule{110mm}{0.4pt}                                                          \\[-1.ex]
+            &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
             &\bf{return} \:  \theta_t                                                     \\[-1.ex]
-            &\rule{110mm}{0.4pt}                                                          \\[-1.ex]
+            &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
        \end{aligned}
 
     .. warning::