mindspore 2.3.0__cp39-cp39-win_amd64.whl → 2.4.1__cp39-cp39-win_amd64.whl

mindspore/experimental/llm_boost/atb/llama_boost.py

@@ -0,0 +1,115 @@
+# 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.
+# ============================================================================
+"""llm boost"""
+import json
+import mindspore.common.dtype as mstype
+from mindspore.experimental.llm_boost.atb.boost_base import AtbBoostBase
+from mindspore._c_expression import LlmBoostBinder
+from mindspore.experimental.llm_boost.register import LlmBoostRegister, LlmBoostType
+
+
+@LlmBoostRegister.register(LlmBoostType.BUILDIN, "Llama")
+class LlamaBoost(AtbBoostBase):
+    """LlamaBoost class"""
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.in_tensor_length = 13
+        self.acl_encoder_operation_inputs = [None] * self.in_tensor_length
+        self.acl_decoder_operation_inputs = [None] * self.in_tensor_length
+        self.atb_encoder_operation = LlmBoostBinder(
+            "ATB", "llama_parallel_DecoderModel")
+        self.atb_decoder_operation = LlmBoostBinder(
+            "ATB", "llama_parallel_DecoderModel")
+
+    def init(self):
+        """set param"""
+        coder_param = {
+            "rmsNormEps": self.config.rms_norm_eps,
+            "numAttentionHeadsPerRank": self.config.num_heads // self.device_num,
+            "hiddenSizePerAttentionHead": self.head_dim,
+            "numHiddenLayers": self.num_layers,
+            "numKeyValueHeadsPerRank": self.n_kv_heads // self.device_num,
+            "skipWordEmbedding": False,
+            "isFA": False,
+            "isBF16": self.dtype == mstype.bfloat16,
+            "packQuantType": [[1, 1] for _ in range(self.num_layers)],
+            "linearQuantType": [[0, -1, -1, 0, 0, -1, 0] for _ in range(self.num_layers)],
+            "linearTransposeType": [[1, -1, -1, 1, 1, -1, 1] for i in range(self.num_layers)],
+            "isEmbeddingParallel": False,
+            "isLmHeadParallel": not self.config.parallel_config.vocab_emb_dp,
+            "lmHeadTransposeType": 1,
+            "supportSwiGLU": True,
+            "kvQuant": self.kv_quant is not None,
+            "rank": self.rank_id,
+            "worldSize": self.device_num,
+            "backend": "lccl",
+            "rankTableFile": "",
+            "positionEmbeddingType": self.position_embedding_type,
+            "hiddenSize": self.config.hidden_size,
+            "gemma": False,
+            "enableAddNorm": True,
+            "supportCompressHead": False,
+        }
+        encoder_param = {
+            **coder_param, "isPrefill": True,
+            "supportLcoc": True,
+            "supportSpeculate": False,
+            "skipWordEmbedding": False
+        }
+        decoder_param = {
+            **coder_param, "isPrefill": False, "supportLcoc": False,
+            "supportSpeculate": False
+        }
+        self.atb_encoder_operation.init(json.dumps({**encoder_param}))
+        self.atb_decoder_operation.init(json.dumps({**decoder_param}))
+
+    def _prepare_inputs(
+            self,
+            prefill=None,
+            input_ids=None,
+            position_ids=None,
+            cos_embed=None,
+            sin_embed=None,
+            attention_mask=None,
+            block_tables=None,
+            slots=None,
+            input_lengths=None,
+            lm_head_indices=None,
+            seqLen=None,
+            **kwargs
+    ):
+        """prepare inputs"""
+        self.acl_param = json.dumps({
+            "seqLen": seqLen,
+        })
+        self.acl_decoder_operation_inputs[0] = self.cast(
+            input_ids, mstype.int64)
+        self.acl_decoder_operation_inputs[1] = self.placeholder
+        self.acl_decoder_operation_inputs[2] = self.cast(
+            position_ids, mstype.int32)
+        self.acl_decoder_operation_inputs[3] = cos_embed
+        self.acl_decoder_operation_inputs[4] = sin_embed
+        self.acl_decoder_operation_inputs[5] = attention_mask
+        self.acl_decoder_operation_inputs[6] = block_tables
+        self.acl_decoder_operation_inputs[7] = slots
+        self.acl_decoder_operation_inputs[8] = self.placeholder
+        self.acl_decoder_operation_inputs[9] = self.placeholder
+        self.acl_decoder_operation_inputs[10] = self.placeholder
+        self.acl_decoder_operation_inputs[11] = self.cast(
+            input_lengths, mstype.int32)
+        self.acl_decoder_operation_inputs[12] = self.cast(
+            lm_head_indices, mstype.int64)
+        return self.acl_decoder_operation_inputs, self.acl_param

mindspore/experimental/llm_boost/atb/qwen_boost.py

@@ -0,0 +1,101 @@
+# 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.
+# ============================================================================
+"""llm boost"""
+import json
+import mindspore.common.dtype as mstype
+from mindspore.experimental.llm_boost.atb.boost_base import AtbBoostBase
+from mindspore._c_expression import LlmBoostBinder
+from mindspore.experimental.llm_boost.register import LlmBoostRegister, LlmBoostType
+
+
+@LlmBoostRegister.register(LlmBoostType.BUILDIN, "Qwen")
+class QwenBoost(AtbBoostBase):
+    """QwenBoost class"""
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.in_tensor_length = 12
+        self.acl_encoder_operation_inputs = [None] * self.in_tensor_length
+        self.acl_decoder_operation_inputs = [None] * self.in_tensor_length
+        self.atb_encoder_operation = LlmBoostBinder(
+            "ATB", "qwen_DecoderModel")
+        self.atb_decoder_operation = LlmBoostBinder(
+            "ATB", "qwen_DecoderModel")
+
+    def init(self):
+        """set param"""
+        param_dict = {
+            "isFA": False,
+            "isBF16": self.dtype == mstype.bfloat16,
+            "withEmbedding": True,
+            "isEmbeddingParallel": True,
+            "isLmHeadParallel": True,
+            "linearTransposeType": [[1, -1, -1, 1, 1, -1, 1] for i in range(self.num_layers)],
+            "lmHeadTransposeType": 1,
+            "supportSwiGLU": not self.need_nz,
+            "rmsNormEps": self.config.rms_norm_eps,
+            "numAttentionHeadsPerRank": self.config.num_heads // self.device_num,
+            "hiddenSizePerAttentionHead": self.head_dim,
+            "numHiddenLayers": self.num_layers,
+            "numKeyValueHeadsPerRank": self.n_kv_heads // self.device_num,
+            "rank": self.rank_id,
+            "worldSize": self.device_num,
+            "backend": "lccl",
+            "packQuantType": [[1, 1] for _ in range(self.num_layers)],
+            "linearQuantType": [[0, -1, -1, 0, 0, -1, 0] for _ in range(self.num_layers)],
+            "kvQuant": self.kv_quant is not None,
+        }
+        encoder_param = {**param_dict, "isPrefill": True, "supportLcoc": False}
+        decoder_param = {**param_dict, "isPrefill": False,
+                         "supportLcoc": False, "supportSpeculate": False}
+        self.atb_encoder_operation.init(json.dumps({**encoder_param}))
+        self.atb_decoder_operation.init(json.dumps({**decoder_param}))
+
+    def _prepare_inputs(
+            self,
+            prefill=None,
+            input_ids=None,
+            position_ids=None,
+            cos_embed=None,
+            sin_embed=None,
+            attention_mask=None,
+            block_tables=None,
+            slots=None,
+            input_lengths=None,
+            lm_head_indices=None,
+            seqLen=None,
+            **kwargs
+    ):
+        """prepare inputs"""
+        self.acl_param = json.dumps({
+            "seqLen": seqLen,
+        })
+        self.acl_decoder_operation_inputs[0] = self.cast(
+            input_ids, mstype.int64)
+        self.acl_decoder_operation_inputs[1] = self.cast(
+            position_ids, mstype.int32)
+        self.acl_decoder_operation_inputs[2] = cos_embed
+        self.acl_decoder_operation_inputs[3] = sin_embed
+        self.acl_decoder_operation_inputs[4] = attention_mask
+        self.acl_decoder_operation_inputs[5] = block_tables
+        self.acl_decoder_operation_inputs[6] = slots
+        self.acl_decoder_operation_inputs[7] = self.placeholder
+        self.acl_decoder_operation_inputs[8] = self.placeholder
+        self.acl_decoder_operation_inputs[9] = self.cast(
+            input_lengths, mstype.int32)
+        self.acl_decoder_operation_inputs[10] = self.cast(
+            lm_head_indices, mstype.int64)
+        self.acl_decoder_operation_inputs[11] = self.placeholder
+        return self.acl_decoder_operation_inputs, self.acl_param

mindspore/experimental/llm_boost/register.py

@@ -0,0 +1,129 @@
+# 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.
+# ============================================================================
+"""LlmBoostRegister"""
+import inspect
+
+
+class LlmBoostType:
+    """Class module type for vision pretrain"""
+
+    def __init__(self):
+        pass
+
+    BUILDIN = 'BuildIn'
+
+
+class LlmBoostRegister:
+    """
+    Module class factory.
+    """
+
+    def __init__(self):
+        pass
+
+    registry = {}
+
+    @classmethod
+    def register(cls, boost_type=LlmBoostType.BUILDIN, alias=None):
+        """Register class into registry
+        Args:
+            boost_type:
+                boost type name, default LlmBoostType.BUILDIN
+            alias (str) : model_name
+
+        Returns:
+            wrapper
+        """
+
+        def wrapper(register_class):
+            """Register-Class with wrapper function.
+
+            Args:
+                register_class : class need to register
+
+            Returns:
+                wrapper of register_class
+            """
+            model_name = alias if alias is not None else register_class.__name__
+            if boost_type not in cls.registry:
+                cls.registry[boost_type] = {model_name: register_class}
+            else:
+                cls.registry[boost_type][model_name] = register_class
+            return register_class
+
+        return wrapper
+
+    @classmethod
+    def is_exist(cls, boost_type, model_name=None):
+        """Determine whether class name is in the current type group.
+
+        Args:
+            boost_type : Module type
+            model_name : model name
+
+        Returns:
+            True/False
+        """
+        if not model_name:
+            return boost_type in cls.registry
+        registered = boost_type in cls.registry and model_name in cls.registry.get(
+            boost_type)
+        return registered
+
+    @classmethod
+    def get_cls(cls, boost_type, model_name=None):
+        """Get class
+
+        Args:
+            boost_type : Module type
+            model_name : model name
+
+        Returns:
+            register_class
+        """
+        if not cls.is_exist(boost_type, model_name):
+            raise ValueError("Can't find class type {} class name {} \
+            in class registry".format(boost_type, model_name))
+
+        if not model_name:
+            raise ValueError(
+                "Can't find model. model name = {}".format(model_name))
+        register_class = cls.registry.get(boost_type).get(model_name)
+        return register_class
+
+    @classmethod
+    def get_instance(cls, boost_type=LlmBoostType.BUILDIN, model_name=None, **kwargs):
+        """Get instance.
+        Args:
+            boost_type : module type
+            model_name : model type
+        Returns:
+            object : The constructed object
+        """
+        if model_name is None:
+            raise ValueError("Class name cannot be None.")
+
+        if isinstance(model_name, str):
+            obj_cls = cls.get_cls(boost_type, model_name)
+        elif inspect.isclass(model_name):
+            obj_cls = model_name
+        else:
+            raise ValueError("Can't find boost type {} model name {} \
+            in class registry.".format(boost_type, model_name))
+
+        try:
+            return obj_cls(**kwargs)
+        except Exception as e:
+            raise type(e)('{}: {}'.format(obj_cls.__name__, e))

mindspore/experimental/llm_boost/utils.py

@@ -0,0 +1,31 @@
+# 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.
+# ============================================================================
+"""LlmBoostRegister"""
+import os
+from mindspore.communication import get_group_size, get_rank
+
+
+def get_real_rank():
+    try:
+        return get_rank()
+    except RuntimeError:
+        return int(os.getenv("RANK_ID", "0"))
+
+
+def get_real_group_size():
+    try:
+        return get_group_size()
+    except RuntimeError:
+        return int(os.getenv("RANK_SIZE", "1"))

mindspore/experimental/optim/adamw.py

@@ -20,10 +20,13 @@ from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
 import mindspore.common.dtype as mstype
 from mindspore.experimental.optim.optimizer import Optimizer
+from mindspore import _checkparam as validator
+from mindspore.ops import auto_generate as gen
 from mindspore import ops
 from mindspore import jit
 
 _adamw_opt = C.MultitypeFuncGraph("adamw_opt")
+_speed_adamw_opt = C.MultitypeFuncGraph("speed_adamw_opt")
 
 op_mul = P.Mul()
 op_pow = P.Pow()
@@ -32,6 +35,30 @@ op_maximum = P.Maximum()
 hyper_map = C.HyperMap()
 
 
+@_speed_adamw_opt.register("Function", "Float", "Float", "Tensor", "Float", "Float", "Bool", "Bool", "Tensor", "Tensor",
+                           "Tensor", "Tensor", "Tensor", "Tensor")
+def _run_speed_adamw_opt(opt, beta1, beta2, lr, eps, weight_decay, amsgrad, maximize, bias_correction1,
+                         bias_correction2, parameters, grads, exp_avg, exp_avg_sq):
+    """Apply adamw optimizer to the weight parameter."""
+    success = True
+    opt(parameters, exp_avg, exp_avg_sq, bias_correction1, bias_correction2, lr, weight_decay, beta1, beta2, eps,
+        grads, None, amsgrad, maximize)
+    return success
+
+
+def _check_param_value(betas, eps, weight_decay, lr, amsgrad, maximize, prim_name):
+    """Check the type of inputs."""
+    validator.check_value_type('betas', betas, [tuple], prim_name)
+    validator.check("betas size", len(betas), "", [2], validator.IN, prim_name)
+    validator.check_value_type("betas[0]", betas[0], [float], prim_name)
+    validator.check_value_type("betas[1]", betas[1], [float], prim_name)
+    validator.check_value_type("eps", eps, [float], prim_name)
+    validator.check_value_type("weight_decay", weight_decay, [float], prim_name)
+    validator.check_value_type("lr", lr, [float], prim_name)
+    validator.check_value_type("amsgrad", amsgrad, [bool], prim_name)
+    validator.check_value_type("maximize", maximize, [bool], prim_name)
+
+
 @jit
 def prepare_func(lr, weight_decay, state_step, beta1, beta2):
     weight_decay_new = 1 - lr * weight_decay
@@ -203,3 +230,61 @@ class AdamW(Optimizer):
                                 grads, start_id, end_id)
 
         return True
+
+
+class SpeedAdamW(Optimizer):
+    r"""
+    Implements Adam Weight Decay algorithm.
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=1e-2, amsgrad=False, *, maximize=False):
+        _check_param_value(betas, eps, weight_decay, lr, amsgrad, maximize, self.cls_name)
+        if lr < 0.0:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if eps < 0.0:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        if weight_decay < 0.0:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                        weight_decay=weight_decay, amsgrad=amsgrad,
+                        maximize=maximize)
+        self.max_v_group = True
+        super(SpeedAdamW, self).__init__(params, defaults)
+
+        self.exp_avg = self.parameters.clone(prefix="exp_avg", init='zeros')
+        self.exp_avg_sq = self.parameters.clone(prefix="exp_avg_sq", init='zeros')
+        self.state_step = Parameter(Tensor([0], mstype.float32), "state_step")
+        self.increase_tensor = Tensor(1, mstype.float32)
+        self.assignadd = P.AssignAdd()
+        self.adamw_opt = gen.ApplyAdamW()
+
+    def construct(self, gradients):
+        self.assignadd(self.state_step, self.increase_tensor)
+        for group_id, group in enumerate(self.param_groups):
+            beta1, beta2 = group['betas']
+            maximize = group.get("maximize")
+            start_id = self.group_start_id[group_id]
+            end_id = self.group_start_id[group_id + 1]
+            lr = group.get("lr")
+            grads = tuple(gradients[start_id: end_id])
+
+            bias_correction1 = float(beta1) ** (float(self.state_step) - 1.0)
+            bias_correction2 = float(beta2) ** (float(self.state_step) - 1.0)
+
+            # 当前 ApplyAdamW 仅支持 amsgrad 为 False
+            if group.get("amsgrad"):
+                raise ValueError("For SpeedAdamW, the value of amsgrad can only be False.")
+
+            self.hyper_map(F.partial(_speed_adamw_opt, self.adamw_opt, beta1, beta2, lr,
+                                     group.get("eps"), group.get("weight_decay"),
+                                     group.get("amsgrad"), maximize, bias_correction1, bias_correction2),
+                           self.parameters[start_id: end_id], grads, self.exp_avg[start_id: end_id],
+                           self.exp_avg_sq[start_id: end_id])
+
+        return True

mindspore/experimental/optim/optimizer.py

@@ -143,6 +143,9 @@ class Optimizer(Cell):
         self.lrs.append(lr)
         param_group["lr"] = lr
         param_group["weight_decay"] = weight_decay
+        if "amsgrad" in param_group and param_group.get("amsgrad") and hasattr(self, 'max_v_group'):
+            param_items = ParameterTuple(tuple(param_group.get("params")))
+            param_group["max_exp_avg_sq"] = param_items.clone(prefix="max_exp_avg_sq", init='zeros')
         self.param_groups.append(param_group)
         self.group_start_id.append(self.group_start_id[-1] + len(param_group.get("params")))
 

mindspore/hal/__init__.py

@@ -22,7 +22,7 @@ resources at the Python layer. Currently, these interfaces take effect only in P
 from mindspore.hal.device import is_initialized, is_available, device_count, get_device_capability,\
                                  get_device_properties, get_device_name, get_arch_list
 from mindspore.hal.stream import Stream, synchronize, set_cur_stream, current_stream, default_stream,\
-                                 StreamCtx
+                                 communication_stream, StreamCtx
 from mindspore.hal.event import Event
 from mindspore.hal.memory import memory_stats, memory_reserved, max_memory_reserved, empty_cache,\
                                  reset_peak_memory_stats, memory_summary, memory_allocated,\
@@ -31,8 +31,8 @@ from mindspore.hal.memory import memory_stats, memory_reserved, max_memory_reser
 __all__ = [
     "is_initialized", "is_available", "device_count", "get_device_capability",
     "get_device_properties", "get_device_name", "get_arch_list",
-    "Event", "Stream", "synchronize", "set_cur_stream", "current_stream", "default_stream", "StreamCtx",
-    "memory_stats", "memory_reserved", "max_memory_reserved", "empty_cache", "reset_peak_memory_stats",
+    "Event", "communication_stream", "Stream", "synchronize", "set_cur_stream", "current_stream", "default_stream",
+    "StreamCtx", "memory_stats", "memory_reserved", "max_memory_reserved", "empty_cache", "reset_peak_memory_stats",
     "memory_summary", "memory_allocated", "max_memory_allocated", "reset_max_memory_reserved",
     "reset_max_memory_allocated"
 ]