configuration_llama.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. # Copyright 2022 EleutherAI and the 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. """LLaMA model configuration"""
  20. from huggingface_hub.dataclasses import strict
  21. from ...configuration_utils import PreTrainedConfig
  22. from ...modeling_rope_utils import RopeParameters
  23. from ...utils import auto_docstring
  24. from ...utils.type_validators import interval
  25. @auto_docstring(checkpoint="meta-llama/Llama-2-7b-hf")
  26. @strict
  27. class LlamaConfig(PreTrainedConfig):
  28. r"""
  29. ```python
  30. >>> from transformers import LlamaModel, LlamaConfig
  31. >>> # Initializing a LLaMA llama-7b style configuration
  32. >>> configuration = LlamaConfig()
  33. >>> # Initializing a model from the llama-7b style configuration
  34. >>> model = LlamaModel(configuration)
  35. >>> # Accessing the model configuration
  36. >>> configuration = model.config
  37. ```"""
  38. model_type = "llama"
  39. keys_to_ignore_at_inference = ["past_key_values"]
  40. # Default tensor parallel plan for base model `LlamaModel`
  41. base_model_tp_plan = {
  42. "layers.*.self_attn.q_proj": "colwise",
  43. "layers.*.self_attn.k_proj": "colwise",
  44. "layers.*.self_attn.v_proj": "colwise",
  45. "layers.*.self_attn.o_proj": "rowwise",
  46. "layers.*.mlp.gate_proj": "colwise",
  47. "layers.*.mlp.up_proj": "colwise",
  48. "layers.*.mlp.down_proj": "rowwise",
  49. }
  50. base_model_pp_plan = {
  51. "embed_tokens": (["input_ids"], ["inputs_embeds"]),
  52. "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
  53. "norm": (["hidden_states"], ["hidden_states"]),
  54. }
  55. vocab_size: int = 32000
  56. hidden_size: int = 4096
  57. intermediate_size: int = 11008
  58. num_hidden_layers: int = 32
  59. num_attention_heads: int = 32
  60. num_key_value_heads: int | None = None
  61. hidden_act: str = "silu"
  62. max_position_embeddings: int = 2048
  63. initializer_range: float = interval(min=0.0, max=1.0)(default=0.02)
  64. rms_norm_eps: float = 1e-6
  65. use_cache: bool = True
  66. pad_token_id: int | None = None
  67. bos_token_id: int | None = 1
  68. eos_token_id: int | list[int] | None = 2
  69. pretraining_tp: int | None = 1
  70. tie_word_embeddings: bool = False
  71. rope_parameters: RopeParameters | dict | None = None
  72. attention_bias: bool = False
  73. attention_dropout: int | float | None = 0.0
  74. mlp_bias: bool = False
  75. head_dim: int | None = None
  76. def __post_init__(self, **kwargs):
  77. if self.head_dim is None:
  78. self.head_dim = self.hidden_size // self.num_attention_heads
  79. if self.num_key_value_heads is None:
  80. self.num_key_value_heads = self.num_attention_heads
  81. super().__post_init__(**kwargs)
  82. def validate_architecture(self):
  83. """Part of `@strict`-powered validation. Validates the architecture of the config."""
  84. if self.hidden_size % self.num_attention_heads != 0:
  85. raise ValueError(
  86. f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention "
  87. f"heads ({self.num_attention_heads})."
  88. )
  89. __all__ = ["LlamaConfig"]