modular_youtu.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  1. # Copyright 2026 the Tencent and HuggingFace Inc. team. All rights reserved.
  2. #
  3. # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
  4. # and OPT implementations in this library. It has been modified from its
  5. # original forms to accommodate minor architectural differences compared
  6. # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
  7. #
  8. # Licensed under the Apache License, Version 2.0 (the "License");
  9. # you may not use this file except in compliance with the License.
  10. # You may obtain a copy of the License at
  11. #
  12. # http://www.apache.org/licenses/LICENSE-2.0
  13. #
  14. # Unless required by applicable law or agreed to in writing, software
  15. # distributed under the License is distributed on an "AS IS" BASIS,
  16. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  17. # See the License for the specific language governing permissions and
  18. # limitations under the License.
  19. import torch
  20. from huggingface_hub.dataclasses import strict
  21. from torch import nn
  22. from ... import initialization as init
  23. from ...modeling_utils import PreTrainedModel
  24. from ...utils import auto_docstring, logging
  25. from ..deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config
  26. from ..deepseek_v3.modeling_deepseek_v3 import DeepseekV3Attention
  27. from ..llama.modeling_llama import (
  28. LlamaDecoderLayer,
  29. LlamaForCausalLM,
  30. LlamaModel,
  31. LlamaPreTrainedModel,
  32. LlamaRMSNorm,
  33. LlamaRotaryEmbedding,
  34. )
  35. from ..qwen3.modeling_qwen3 import Qwen3MLP
  36. logger = logging.get_logger(__name__)
  37. @auto_docstring(checkpoint="tencent/Youtu-LLM-2B")
  38. @strict
  39. class YoutuConfig(DeepseekV3Config):
  40. r"""
  41. rope_interleave (`bool`, *optional*, defaults to `True`):
  42. Whether to interleave the rotary position embeddings.
  43. embedding_initializer_range (`float`, *optional*):
  44. The standard deviation of the truncated_normal_initializer for initializing all embedding matrices.
  45. ```python
  46. >>> from transformers import YoutuModel, YoutuConfig
  47. >>> # Initializing a Youtu-LLM-2B style configuration
  48. >>> configuration = YoutuConfig()
  49. >>> # Accessing the model configuration
  50. >>> configuration = model.config
  51. ```"""
  52. model_type = "youtu"
  53. base_model_tp_plan = {
  54. "layers.*.mlp.gate_proj": "colwise",
  55. "layers.*.mlp.up_proj": "colwise",
  56. "layers.*.mlp.down_proj": "rowwise",
  57. }
  58. attribute_map = {}
  59. vocab_size: int = 128256
  60. hidden_size: int = 2048
  61. intermediate_size: int = 6144
  62. num_hidden_layers: int = 32
  63. num_attention_heads: int = 16
  64. num_key_value_heads: int = 16
  65. max_position_embeddings: int = 131072
  66. initializer_range: float | None = None
  67. embedding_initializer_range: float | None = None
  68. pad_token_id: int | None = None
  69. bos_token_id: int | None = 128000
  70. eos_token_id: int | list[int] | None = 128001
  71. tie_word_embeddings: bool = True
  72. # remove unused attribute
  73. n_shared_experts = AttributeError()
  74. n_routed_experts = AttributeError()
  75. routed_scaling_factor = AttributeError()
  76. n_group = AttributeError()
  77. topk_group = AttributeError()
  78. num_experts_per_tok = AttributeError()
  79. first_k_dense_replace = AttributeError()
  80. norm_topk_prob = AttributeError()
  81. pretraining_tp = AttributeError()
  82. moe_intermediate_size = AttributeError()
  83. def __post_init__(self, **kwargs):
  84. if self.initializer_range is None:
  85. if self.hidden_size != 0:
  86. self.initializer_range = 2.0 / (5.0 * self.hidden_size) ** 0.5
  87. else:
  88. self.initializer_range = 0.02
  89. self.embedding_initializer_range = self.embedding_initializer_range or 2.0 * self.initializer_range
  90. super().__post_init__(**kwargs)
  91. def convert_rope_params_to_dict(self, **kwargs):
  92. raise AttributeError("Not overwritten for the Youtu model!")
  93. class YoutuRMSNorm(LlamaRMSNorm):
  94. pass
  95. class YoutuRotaryEmbedding(LlamaRotaryEmbedding):
  96. pass
  97. class YoutuMLP(Qwen3MLP):
  98. pass
  99. class YoutuAttention(DeepseekV3Attention):
  100. pass
  101. class YoutuDecoderLayer(LlamaDecoderLayer):
  102. pass
  103. class YoutuPreTrainedModel(LlamaPreTrainedModel, PreTrainedModel):
  104. @torch.no_grad()
  105. def _init_weights(self, module):
  106. PreTrainedModel._init_weights(self, module)
  107. std = getattr(self.config, "initializer_range", 0.02)
  108. embed_std = getattr(self.config, "embedding_initializer_range", 2 * std)
  109. if isinstance(module, nn.Embedding):
  110. init.normal_(module.weight, mean=0.0, std=embed_std)
  111. if module.padding_idx is not None:
  112. init.zeros_(module.weight.data[module.padding_idx])
  113. class YoutuModel(LlamaModel):
  114. pass
  115. class YoutuForCausalLM(LlamaForCausalLM):
  116. pass
  117. __all__ = [
  118. "YoutuConfig",
  119. "YoutuPreTrainedModel",
  120. "YoutuModel",
  121. "YoutuForCausalLM",
  122. ]