diff --git a/configs/telechat3/predict_telechat3_36b.yaml b/configs/telechat3/predict_telechat3_36b.yaml new file mode 100644 index 0000000000000000000000000000000000000000..53b44e1d3ac0152acf8fbd8388ae3dd86f68f5f5 --- /dev/null +++ b/configs/telechat3/predict_telechat3_36b.yaml @@ -0,0 +1,55 @@ +seed: 0 +output_dir: './output' # path to save checkpoint/strategy +load_checkpoint: '' +use_parallel: True +run_mode: 'predict' +use_legacy: False +load_ckpt_format: 'safetensors' + +trainer: + type: CausalLanguageModelingTrainer + model_name: 'qwen3' + +# default parallel of device num = 8 for Atlas 800T A2 +parallel_config: + data_parallel: 1 + model_parallel: 1 +# HuggingFace file directory +pretrained_model_dir: '/path/hf_dir' +model: + model_config: + seq_length: 32768 + compute_dtype: "bfloat16" + layernorm_compute_dtype: "float32" + softmax_compute_dtype: "float32" + rotary_dtype: "float32" + params_dtype: "bfloat16" + block_size: 128 + num_blocks: 1024 + +generation_config: + do_sample: False + max_length: 32768 + max_new_tokens: 512 + bos_token_id: 1 + eos_token_id: [2] + pad_token_id: 3 + top_k: 1 + top_p: 1.0 + repetition_penalty: 1.05 + temperature: 1.0 + +# mindspore context init config +context: + mode: 0 #0--Graph Mode; 1--Pynative Mode + enable_graph_kernel: False + ascend_config: + precision_mode: "must_keep_origin_dtype" + max_device_memory: "59GB" + save_graphs: False + save_graphs_path: "./graph" + +# parallel context config +parallel: + parallel_mode: "MANUAL_PARALLEL" + enable_alltoall: False \ No newline at end of file diff --git a/mindformers/models/telechat3/configuration_telechat3.py b/mindformers/models/telechat3/configuration_telechat3.py index 300649543b0173f6c608d2b8734726b5e11a5a6d..cec7cb6cbdb508f44ef96a9f38356646fb22181a 100644 --- a/mindformers/models/telechat3/configuration_telechat3.py +++ b/mindformers/models/telechat3/configuration_telechat3.py @@ -37,6 +37,13 @@ from mindformers.models.model_config_utils import ( @MindFormerRegister.register(MindFormerModuleType.CONFIG, legacy=False, search_names=['telechat3', 'telechat3_mtp']) class TeleChat3Config(PretrainedConfig): + def __new__(cls, **kwargs): + is_moe = kwargs.get("n_routed_experts", None) + if is_moe is None: + return TeleChat3DenseConfig(**kwargs) + return TeleChat3MoeConfig(**kwargs) + +class TeleChat3MoeConfig(PretrainedConfig): r""" This is the configuration class to store the configuration of a [`TeleChat3Model`]. It is used to instantiate an TeleChat3 model according to the specified arguments, defining the model architecture. @@ -240,3 +247,116 @@ class TeleChat3Config(PretrainedConfig): tie_word_embeddings=tie_word_embeddings, **kwargs, ) + +class TeleChat3DenseConfig(PretrainedConfig): + model_type = "telechat3" + keys_to_ignore_at_inference = ["past_key_values"] + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + @register_mf_model_parameter( + mf_model_kwargs=MFModelConfig( + pad_token_id=3, + block_size=32, + num_blocks=1024, + normalization='RMSNorm', + add_bias_linear=False, + gated_linear_unit=True, + hidden_act='silu' + )) + @ignore_and_delete_parameter(extra_ignore_param=[ + ('max_window_layers', NotSupportedInfo.useless), + ('sliding_window', NotSupportedInfo.useless), + ('use_sliding_window', NotSupportedInfo.useless), + ('layer_types', NotSupportedInfo.useless), + ('logn', NotSupportedInfo.useless), + ('training_seqlen', NotSupportedInfo.useless), + ('embed_layernorm', NotSupportedInfo.useless), + ('unk_token_id', NotSupportedInfo.useless), + ('flash_attn', NotSupportedInfo.useless), + ('base_seqlen', NotSupportedInfo.useless), + ('mlp_bias', NotSupportedInfo.useless), + ('pretraining_tp', NotSupportedInfo.useless), + ('share_attention', NotSupportedInfo.useless), + ('share_ffn', NotSupportedInfo.useless), + ]) + def __init__( + self, + attention_bias=False, + attention_dropout=0.0, + bos_token_id=1, + eos_token_id=2, + head_dim=128, + hidden_act="silu", + hidden_size=6144, + initializer_range=0.0048, + intermediate_size=24576, + max_position_embeddings=2048, + mlp_bias=False, + model_type="telechat3", + num_attention_heads=48, + num_hidden_layers=64, + num_key_value_heads=None, + original_max_position_embeddings=8192, + pad_token_id=None, + pretraining_tp=1, + rms_norm_eps=1e-5, + rope_scaling=None, + rope_theta=1000000.0, + tie_word_embeddings=False, + use_cache=True, + vocab_size=131072, + **kwargs, + ): + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.hidden_size = hidden_size + self.hidden_act = hidden_act + self.intermediate_size = intermediate_size + self.mlp_bias = mlp_bias + self.max_position_embeddings = max_position_embeddings + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + + self.initializer_range = initializer_range + + self.pretraining_tp = pretraining_tp + self.rms_norm_eps = rms_norm_eps + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.use_cache = use_cache + self.vocab_size = vocab_size + + if head_dim is not None and head_dim != self.hidden_size // self.num_attention_heads: + raise ValueError("head_dim != hidden_size//num_attention_head.Please check the config.") + self.head_dim = head_dim if head_dim is not None else self.hidden_size // self.num_attention_heads + + # Validate the correctness of rotary position embeddings parameters + # BC: if there is a 'type' field, copy it it to 'rope_type'. + if self.rope_scaling is not None and "type" in self.rope_scaling: + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) \ No newline at end of file diff --git a/mindformers/models/telechat3/modeling_telechat3.py b/mindformers/models/telechat3/modeling_telechat3.py index 033cf9745d31857210f3a7b7981e78e623f1850b..b1d00d3b05c1eaa113af649f67d8237d5e5fb1a0 100644 --- a/mindformers/models/telechat3/modeling_telechat3.py +++ b/mindformers/models/telechat3/modeling_telechat3.py @@ -13,10 +13,12 @@ # limitations under the License. # ============================================================================ """Telechat3 Model.""" +import os from mindformers.tools.register import MindFormerRegister, MindFormerModuleType from .configuration_telechat3 import TeleChat3Config +from .modeling_telechat3_infer import InferenceTelechat3ForCausalLM from .modeling_telechat3_train import TrainingTeleChat3ForCausalLM @@ -34,4 +36,6 @@ class TeleChat3ForCausalLM: def __new__(cls, config: TeleChat3Config, *args, **kwargs): # get run mode to init different model. # we can use train mode to do train task. + if os.environ.get("RUN_MODE") == "predict": + return InferenceTelechat3ForCausalLM(config=config) return TrainingTeleChat3ForCausalLM(config=config) diff --git a/mindformers/models/telechat3/modeling_telechat3_infer.py b/mindformers/models/telechat3/modeling_telechat3_infer.py new file mode 100644 index 0000000000000000000000000000000000000000..881486718b49bdec7e137d78112ddc6c67d8de21 --- /dev/null +++ b/mindformers/models/telechat3/modeling_telechat3_infer.py @@ -0,0 +1,120 @@ +# Copyright 2025 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. +# ============================================================================ +"""Telechat3 models' APIs.""" +__all__ = ['InferenceTelechat3ForCausalLM'] + +from mindformers.models.utils import jit +from mindformers.parallel_core.transformer_config import TransformerConfig +from mindformers.models.telechat3.utils import TeleChat3PreTrainedModel +from mindformers.parallel_core.inference.base_models.gpt.gpt_model import GPTModel +from mindformers.parallel_core.inference.base_models.gpt.gpt_layer_specs import get_gpt_layer_local_spec +from mindformers.parallel_core.inference.model_utils import InferModelMixin +from mindformers.parallel_core.inference.quantization.utils import get_quant_config + +from .configuration_telechat3 import TeleChat3Config + + +class InferenceTelechat3ForCausalLM(TeleChat3PreTrainedModel, InferModelMixin): + r""" + Provide telechat3 model infer through network. + + Args: + config (Telechat3Config): The config of telechat3 model. + + Returns: + output: Tensor, the output of telechat3 decoder layer + + """ + + def __init__(self, config: TeleChat3Config): + super().__init__(config, auto_prefix=False) + self.config = config + config: TransformerConfig = self.convert_to_transformer_config(self.config) + + self.quant_config = get_quant_config(self.config, self.weight_mapping) + self.pad_token_id = self.config.pad_token_id + self.vocab_size = config.vocab_size + self.max_position_embeddings = config.max_position_embeddings + self.compute_dtype = config.compute_dtype + + self.is_prefill = True + self.model = GPTModel(config=config, + transformer_layer_spec=get_gpt_layer_local_spec( + normalization=config.normalization, + use_flash_attention=self.config.use_flash_attention, + qk_layernorm=False, + ), + vocab_size=self.vocab_size, + max_sequence_length=self.max_position_embeddings, + position_embedding_type=config.position_embedding_type, + rotary_base=self.config.rope_theta, + share_embeddings_and_output_weights=self.config.tie_word_embeddings, + pre_process=config.pre_process, + post_process=config.post_process, + quant_config=self.quant_config) + + @jit + def construct(self, input_ids, hidden_states=None, positions=None, batch_valid_length=None, + context_lens_tensor=None, q_seq_lens=None, block_tables=None, slot_mapping=None, + attention_mask=None, attn_metadata=None, attn_padding_idx=None, attn_unpadding_idx=None, + ffn_padding_idx=None, ffn_unpadding_idx=None, key_cache=None, value_cache=None): + """ + model forward. + + Args: + input_ids: input ids. + positions: position ids. + hidden_states: hidden states. + batch_valid_length: actual seq length. + context_lens_tensor: computed key value length. + q_seq_lens: query sequence lengths. + block_tables: Store mapping tables for each sequence. + slot_mapping : Token cache physical slot index. + attention_mask: attention mask used for fa or pa. + attn_metadata: attention metadata + attn_padding_idx: Indices mapping positions in attention output sequence to original token positions, + used for padding attention output to fixed size. + attn_unpadding_idx: Indices mapping valid tokens in padded attention output sequence to + their original positions, used for removing padding in attention output. + ffn_padding_idx: Indices mapping positions in MoE output sequence to flattened valid token positions, + used for padding MoE output to fixed size. + ffn_unpadding_idx: Indices mapping valid tokens in padded MoE output sequence to their original positions, + used for removing padding in MoE output. + key_cache: key cache for incremental inference. + value_cache: value cache for incremental inference. + + Returns: + logits: the output logits. + + """ + logits = self.model( + input_ids=input_ids, + hidden_states=hidden_states, + positions=positions, + batch_valid_length=batch_valid_length, + context_lens_tensor=context_lens_tensor, + q_seq_lens=q_seq_lens, + block_tables=block_tables, + slot_mapping=slot_mapping, + attention_mask=attention_mask, + attn_metadata=attn_metadata, + attn_padding_idx=attn_padding_idx, + attn_unpadding_idx=attn_unpadding_idx, + ffn_padding_idx=ffn_padding_idx, + ffn_unpadding_idx=ffn_unpadding_idx, + key_cache=key_cache, + value_cache=value_cache + ) + return logits \ No newline at end of file diff --git a/mindformers/models/telechat3/utils.py b/mindformers/models/telechat3/utils.py index 313e17cb52487a17aef73d3d5671ed5fe89d779a..92a4ec235942d21ee8ff57a9f67dc56809f025bf 100644 --- a/mindformers/models/telechat3/utils.py +++ b/mindformers/models/telechat3/utils.py @@ -31,6 +31,12 @@ class TeleChat3PreTrainedModel(PreTrainedModel, ModelMixin): weight_mapping = [ ('model.embed_tokens.', 'embedding.word_embeddings.'), + ('.self_attn.q_proj.', '.self_attention.linear_q.'), + ('.self_attn.k_proj.', '.self_attention.linear_k.'), + ('.self_attn.v_proj.', '.self_attention.linear_v.'), + ('.self_attn.o_proj.', '.self_attention.linear_proj.'), + ('.self_attn.q_norm.', '.self_attention.q_layernorm.'), + ('.self_attn.k_norm.', '.self_attention.k_layernorm.'), ('.self_attn.q_a_proj.', '.self_attention.linear_q_down_proj.'), ('.self_attn.q_a_layernorm.', '.self_attention.q_layernorm.'), ('.self_attn.q_b_proj.', '.self_attention.linear_q_up_proj.'), diff --git a/mindformers/parallel_core/inference/base_models/common/embeddings/rope_utils.py b/mindformers/parallel_core/inference/base_models/common/embeddings/rope_utils.py index 04434e66156886cc1d94b0494628a043e3d9d6d6..1715f4965b03755b0ae65f4ed52de8addac223f5 100644 --- a/mindformers/parallel_core/inference/base_models/common/embeddings/rope_utils.py +++ b/mindformers/parallel_core/inference/base_models/common/embeddings/rope_utils.py @@ -27,6 +27,8 @@ from mindformers.parallel_core.inference.base_models.common.embeddings.rotary_po ) from mindformers.parallel_core.inference.base_models.common.embeddings.yarn_rotary_pos_embedding import \ YaRNScalingRotaryEmbedding +from mindformers.parallel_core.inference.base_models.common.embeddings.yarn_rotary_pos_embedding_telechat3 import \ + YaRNScalingRotaryEmbeddingTelechat3 _ROPE_DICT: Dict[Tuple, RotaryEmbedding] = {} @@ -40,7 +42,7 @@ def _get_default(**kwargs): rotary_dtype=kwargs["rotary_dtype"], seq_len_interpolation_factor=kwargs["seq_len_interpolation_factor"], rotary_cos_format=kwargs["rotary_cos_format"], - max_position_embeddings=kwargs["original_max_position_embeddings"], + max_position_embeddings=kwargs["max_position_embeddings"], ) @@ -58,7 +60,6 @@ def _get_llama3(**kwargs): max_position_embeddings=kwargs["max_position_embeddings"], low_freq_factor=getattr(config, "low_freq_factor", 1), high_freq_factor=getattr(config, "low_freq_factor", 8), - orig_max_position=kwargs["original_max_position_embeddings"], ) @@ -73,7 +74,8 @@ def _get_yarn(**kwargs): seq_len_interpolation_factor=kwargs["seq_len_interpolation_factor"], rotary_cos_format=kwargs["rotary_cos_format"], scaling_factor=config.rotary_scaling_factor, - original_max_position_embeddings=kwargs["original_max_position_embeddings"], + max_position_embeddings=kwargs["max_position_embeddings"], + original_max_position_embeddings=config.original_max_position_embeddings, beta_slow=config.beta_slow, beta_fast=config.beta_fast, mscale=config.mscale, @@ -89,10 +91,30 @@ def _get_partialrope(**kwargs): rotary_base=kwargs["rotary_base"], rotary_dtype=kwargs["rotary_dtype"], seq_len_interpolation_factor=kwargs["seq_len_interpolation_factor"], - max_position_embeddings=kwargs["original_max_position_embeddings"], + max_position_embeddings=kwargs["max_position_embeddings"], rotary_cos_format=kwargs["rotary_cos_format"], ) +def _get_telechat3_yarn(**kwargs): + """Instantiate a YaRNScalingRotaryEmbedding object""" + config: MLATransformerConfig = kwargs["config"] + return YaRNScalingRotaryEmbeddingTelechat3( + kv_channels=kwargs["kv_channels"], + rotary_percent=kwargs["rotary_percent"], + rotary_base=kwargs["rotary_base"], + rotary_dtype=kwargs["rotary_dtype"], + seq_len_interpolation_factor=kwargs["seq_len_interpolation_factor"], + rotary_cos_format=kwargs["rotary_cos_format"], + scaling_factor=config.rotary_scaling_factor, + attention_factor=getattr(config, "attention_factor", None), + max_position_embeddings=kwargs["max_position_embeddings"], + original_max_position_embeddings=getattr(config, "original_max_position_embeddings", None), + beta_slow=config.beta_slow, + beta_fast=config.beta_fast, + mscale=getattr(config, "mscale", None), + mscale_all_dim=getattr(config, "mscale_all_dim", None), + ) + # Note: When adding a new RoPE class, add the mapping of function here ROPE_FUNCTION = { @@ -100,6 +122,7 @@ ROPE_FUNCTION = { 'llama': _get_llama3, 'yarn': _get_yarn, 'partial_rope': _get_partialrope, + 'telechat3-yarn': _get_telechat3_yarn, } @@ -111,7 +134,7 @@ def get_rope( rotary_dtype: mstype, seq_len_interpolation_factor: float = None, position_embedding_type: str = 'rope', - original_max_position_embeddings: int = 4096, + max_position_embeddings: int = 4096, rotary_cos_format: str = "rotate_half", **kwargs, ) -> RotaryEmbedding: @@ -123,7 +146,7 @@ def get_rope( f"configuration in the model configuration file config.json.") key = (hidden_dim, rotary_percent, rotary_base, rotary_dtype, position_embedding_type, - original_max_position_embeddings, rotary_cos_format) + tuple(kwargs.values()) + max_position_embeddings, rotary_cos_format) + tuple(kwargs.values()) if key in _ROPE_DICT: return _ROPE_DICT[key] rotary_emb = ROPE_FUNCTION.get(position_embedding_type)( @@ -133,7 +156,7 @@ def get_rope( rotary_base=rotary_base, rotary_dtype=rotary_dtype, seq_len_interpolation_factor=seq_len_interpolation_factor, - original_max_position_embeddings=original_max_position_embeddings, + max_position_embeddings=max_position_embeddings, rotary_cos_format=rotary_cos_format, **kwargs, ) diff --git a/mindformers/parallel_core/inference/base_models/common/embeddings/yarn_rotary_pos_embedding.py b/mindformers/parallel_core/inference/base_models/common/embeddings/yarn_rotary_pos_embedding.py index a68b92b62b8436c8cb52639c1e2efd7f497acfbd..70d34af40aa9977598b7dfba370a7f905604f411 100644 --- a/mindformers/parallel_core/inference/base_models/common/embeddings/yarn_rotary_pos_embedding.py +++ b/mindformers/parallel_core/inference/base_models/common/embeddings/yarn_rotary_pos_embedding.py @@ -59,7 +59,7 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding): rotary_cos_format: str = "rotate_half", rotary_dtype: mstype = mstype.float16, scaling_factor: float = 1.0, - original_max_position_embeddings: int = 4096, + max_position_embeddings: int = 8192, beta_fast: float = 32.0, beta_slow: float = 1.0, mscale: float = 1.0, @@ -68,7 +68,7 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding): attn_factor: float = 1.0 ) -> None: self.scaling_factor = scaling_factor - self.original_max_position_embeddings = original_max_position_embeddings + self.max_position_embeddings = max_position_embeddings self.beta_fast = beta_fast self.beta_slow = beta_slow self.extrapolation_factor = extrapolation_factor @@ -79,7 +79,7 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding): _yarn_get_mscale(self.scaling_factor, float(mscale_all_dim)) * attn_factor) super().__init__(kv_channels, rotary_percent, rotary_interleaved, seq_len_interpolation_factor, - rotary_base, rotary_cos_format, rotary_dtype, original_max_position_embeddings) + rotary_base, rotary_cos_format, rotary_dtype, max_position_embeddings) def _compute_inv_freq(self, base: Union[int, float]) -> Tensor: """Compute the inverse frequency.""" @@ -105,7 +105,7 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding): used to create positional encodings """ inv_freq = self._compute_inv_freq(self.rotary_base) - seq = ops.arange(0, self.max_position_embeddings * self.scaling_factor, 1, dtype=inv_freq.dtype) + offset + seq = ops.arange(0, self.max_position_embeddings, 1, dtype=inv_freq.dtype) + offset freqs = ops.outer(seq, inv_freq) diff --git a/mindformers/parallel_core/inference/base_models/common/embeddings/yarn_rotary_pos_embedding_telechat3.py b/mindformers/parallel_core/inference/base_models/common/embeddings/yarn_rotary_pos_embedding_telechat3.py new file mode 100644 index 0000000000000000000000000000000000000000..b2c25042c2087934507274af9f415dd826cdcaf8 --- /dev/null +++ b/mindformers/parallel_core/inference/base_models/common/embeddings/yarn_rotary_pos_embedding_telechat3.py @@ -0,0 +1,174 @@ +# Copyright 2025 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. +# ============================================================================ +""" +Yarn rotary position embedding for transformer. +""" +__all__ = ['YaRNScalingRotaryEmbeddingTelechat3'] + +from typing import Tuple, Union +import math + +from mindspore import ops +import mindspore.common.dtype as mstype +from mindspore.common.tensor import Tensor + +from mindformers.parallel_core.inference.base_models.common.embeddings.rotary_pos_embedding import RotaryEmbedding + + +class YaRNScalingRotaryEmbeddingTelechat3(RotaryEmbedding): + """RotaryEmbedding extended with YaRN method. + + Args: + kv_channels (int): The number of channels for key and value. + rotary_percent (float): The percentage of the rotary embedding. + rotary_interleaved (bool): Whether to use interleaved rotary position embedding. + seq_len_interpolation_factor (float): The interpolation factor for sequence length. + rotary_base (float): The base for rotary embedding. + rotary_cos_format (str): The mode of ApplyRotaryPosEmb. + rotary_dtype (mstype): The dtype of rotary embeddings. + scaling_factor (float): Base scaling factor for sequence length adjustment. + max_position_embeddings (int): Maximum position embeddings. + original_max_position_embeddings (int): Original maximum position embeddings. + beta_fast (int): High-frequency adjustment cutoff parameter. Larger values + extend high-frequency preservation. Default: 32. + beta_slow (int): Low-frequency adjustment cutoff parameter. Default: 1. + mscale (float): Primary magnitude scaling coefficient. + mscale_all_dim (float): Full-dimensional magnitude scaling coefficient. + extrapolation_factor (float): Weighting factor for extrapolated frequency ranges. + attn_factor (float): Global attention stability coefficient. + + """ + def __init__( + self, + kv_channels: int, + rotary_percent: float = 1.0, + rotary_interleaved: bool = False, + seq_len_interpolation_factor: float = None, + rotary_base: float = 10000.0, + rotary_cos_format: str = "rotate_half", + rotary_dtype: mstype = mstype.float16, + scaling_factor: float = 1.0, + attention_factor: float = None, + max_position_embeddings: int = 4096, + original_max_position_embeddings: int = None, + beta_fast: float = 32.0, + beta_slow: float = 1.0, + mscale: float = None, + mscale_all_dim: float = None, + extrapolation_factor: float = 1.0 + ) -> None: + self.scaling_factor = scaling_factor + self.max_position_embeddings = max_position_embeddings + self.original_max_position_embeddings = original_max_position_embeddings + self.beta_fast = beta_fast + self.beta_slow = beta_slow + self.extrapolation_factor = extrapolation_factor + # Get n-d magnitude scaling corrected for interpolation. + if self.original_max_position_embeddings: + self.scaling_factor = self.max_position_embeddings / self.original_max_position_embeddings + else: + self.original_max_position_embeddings = self.max_position_embeddings + + if attention_factor is None: + if mscale and mscale_all_dim: + self.mscale = float( + _yarn_get_mscale(self.scaling_factor, float(mscale)) / + _yarn_get_mscale(self.scaling_factor, float(mscale_all_dim))) + else: + self.mscale = _yarn_get_mscale(self.scaling_factor) + super().__init__(kv_channels, rotary_percent, rotary_interleaved, seq_len_interpolation_factor, + rotary_base, rotary_cos_format, rotary_dtype, max_position_embeddings) + + def _compute_inv_freq(self, base: Union[int, float]) -> Tensor: + """Compute the inverse frequency.""" + pos_freqs = base**( + ops.arange(0, self.dim, 2, dtype=mstype.float32)[: (self.dim // 2)] / self.dim + ) + inv_freq_extrapolation = 1.0 / pos_freqs + inv_freq_interpolation = 1.0 / (self.scaling_factor * pos_freqs) + + low, high = _yarn_find_correction_range(self.beta_fast, self.beta_slow, + self.dim, base, self.original_max_position_embeddings) + # Get n-d rotational scaling corrected for extrapolation + inv_freq_mask = (1 - _yarn_linear_ramp_mask( + low, high, self.dim // 2, + dtype=mstype.float32)) * self.extrapolation_factor + inv_freq = inv_freq_interpolation * ( + 1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask + return inv_freq + + def get_freqs_non_repeated(self, offset: int = 0) -> Tensor: + """ + Generates matrix of frequencies based on positions in the sequence, + used to create positional encodings + """ + inv_freq = self._compute_inv_freq(self.rotary_base) + seq = ops.arange(0, self.max_position_embeddings, 1, dtype=inv_freq.dtype) + offset + + freqs = ops.outer(seq, inv_freq) + + return freqs + + def _compute_cos_sin_cache(self) -> Tuple[Tensor, Tensor]: + freqs = self.get_freqs_non_repeated() + emb = self.cat((freqs, freqs)) + + cos = self.cast(self.cos(emb) * self.mscale, self.rotary_dtype) + sin = self.cast(self.sin(emb) * self.mscale, self.rotary_dtype) + + return cos, sin + + +def _yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: + """Computes dynamic magnitude scaling coefficient based on scaling factor""" + if scale <= 1: + return 1.0 + return 0.07 * mscale * math.log(scale) + 1.0 + + +def _yarn_find_correction_dim(num_rotations: int, + dim: int, + base: float = 10000, + max_position_embeddings: int = 2048) -> float: + """Inverse dim formula to find dim based on number of rotations""" + return (dim * math.log(max_position_embeddings / + (num_rotations * 2 * math.pi))) / (2 * + math.log(base)) + + +def _yarn_find_correction_range( + low_rot: int, + high_rot: int, + dim: int, + base: float = 10000, + max_position_embeddings: int = 2048) -> Tuple[int, int]: + """Find dim range bounds based on rotations""" + low = math.floor( + _yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)) + high = math.ceil( + _yarn_find_correction_dim(high_rot, dim, base, + max_position_embeddings)) + return max(low, 0), min(high, dim - 1) # Clamp values just in case + + +def _yarn_linear_ramp_mask(low: float, high: float, dim: int, + dtype: mstype) -> Tensor: + """Generates a linear ramp mask for smooth transition.""" + if low == high: + high += 0.001 # Prevent singularity + + linear_func = (ops.arange(dim, dtype=dtype) - low) / (high - low) + ramp_func = ops.clamp(linear_func, 0, 1) + return ramp_func \ No newline at end of file diff --git a/mindformers/parallel_core/inference/base_models/gpt/gpt_model.py b/mindformers/parallel_core/inference/base_models/gpt/gpt_model.py index 71b6498f62d39360c1b09cfcb72f36dcc8b371aa..6b5ec40899da82354edeef542864d09a0ddfdd4a 100644 --- a/mindformers/parallel_core/inference/base_models/gpt/gpt_model.py +++ b/mindformers/parallel_core/inference/base_models/gpt/gpt_model.py @@ -167,7 +167,7 @@ class GPTModel(nn.Cell): rotary_dtype=self.config.rotary_dtype, seq_len_interpolation_factor=self.seq_len_interpolation_factor, position_embedding_type=self.position_embedding_type, - original_max_position_embeddings=self.max_position_embeddings, + max_position_embeddings=self.max_position_embeddings, ) # Transformer diff --git a/mindformers/parallel_core/transformer_config.py b/mindformers/parallel_core/transformer_config.py index 41be3d28cafd6e6a72b20fa0528baae4842ab869..7ec3f14dd9f1dc2e384759a0ed9f1759b3cc9c25 100644 --- a/mindformers/parallel_core/transformer_config.py +++ b/mindformers/parallel_core/transformer_config.py @@ -682,8 +682,6 @@ class TransformerConfig(ModelParallelConfig, MFModelConfig): self.position_embedding_type = (self.rope_scaling.pop("type", None) or self.rope_scaling.pop("rope_type", None)) self.rotary_scaling_factor = self.rope_scaling.pop("factor") - self.max_position_embeddings = self.rope_scaling.pop("original_max_position_embeddings", - None) or self.seq_length for k, v in self.rope_scaling.items(): setattr(self, k, v) del self.rope_scaling diff --git a/mindformers/parallel_core/transformer_config_utils.py b/mindformers/parallel_core/transformer_config_utils.py index 0089ab068e8d3d6b511feb94519ed3f955e0f41c..6428daad8ea197a2bc5108302fa748889c85b646 100644 --- a/mindformers/parallel_core/transformer_config_utils.py +++ b/mindformers/parallel_core/transformer_config_utils.py @@ -182,7 +182,8 @@ PARALLEL_CONFIG_KEY = set(ParallelConfig.keys()) INFER_CONFIG_KEY = set({ "mindformers_version", "rl_config", "checkpoint_name_or_path", "_name_or_path", "type", "model_type", "tokenizer_class", "architectures", "is_encoder_decoder", "is_sample_acceleration", "bos_token_id", "eos_token_id", - "temperature", "repetition_penalty", "max_decode_length", "top_k", "top_p", "do_sample" + "temperature", "repetition_penalty", "max_decode_length", "top_k", "top_p", "do_sample", "generation_config", + "max_length", "max_new_tokens", "post_process" }) DEFAULT_WHITE_KEY.update(PRETRAIN_CONFIG_KEY) DEFAULT_WHITE_KEY.update(PARALLEL_CONFIG_KEY) @@ -389,7 +390,8 @@ COMMON_CONFIG_MAPPING = { ("qk_rope_head_dim", "qk_pos_emb_head_dim"): "qk_pos_emb_head_dim", ("theta", "rope_theta", "rotary_base"): "rotary_base", ("scaling_factor", "factor", "rotary_scaling_factor"): "rotary_scaling_factor", - ("max_position_embeddings", "original_max_position_embeddings"): "max_position_embeddings", + "max_position_embeddings": "max_position_embeddings", + "original_max_position_embeddings" : "original_max_position_embeddings", # not changes "q_lora_rank": "q_lora_rank", "kv_lora_rank": "kv_lora_rank",