configuration_olmo.py 3.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  1. # Copyright 2024 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. """OLMo 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. @auto_docstring(checkpoint="allenai/OLMo-7B-hf")
  25. @strict
  26. class OlmoConfig(PreTrainedConfig):
  27. r"""
  28. clip_qkv (`float`, *optional*):
  29. If not `None`, elements of query, key and value attention states are clipped so that their
  30. absolute value does not exceed this value.
  31. ```python
  32. >>> from transformers import OlmoModel, OlmoConfig
  33. >>> # Initializing a OLMo 7B style configuration
  34. >>> configuration = OlmoConfig()
  35. >>> # Initializing a model from the OLMo 7B style configuration
  36. >>> model = OlmoModel(configuration)
  37. >>> # Accessing the model configuration
  38. >>> configuration = model.config
  39. ```
  40. """
  41. model_type = "olmo"
  42. keys_to_ignore_at_inference = ["past_key_values"]
  43. base_model_tp_plan = {
  44. "layers.*.self_attn.q_proj": "colwise",
  45. "layers.*.self_attn.k_proj": "colwise",
  46. "layers.*.self_attn.v_proj": "colwise",
  47. "layers.*.self_attn.o_proj": "rowwise",
  48. "layers.*.mlp.gate_proj": "colwise",
  49. "layers.*.mlp.up_proj": "colwise",
  50. "layers.*.mlp.down_proj": "rowwise",
  51. }
  52. base_model_pp_plan = {
  53. "embed_tokens": (["input_ids"], ["inputs_embeds"]),
  54. "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
  55. "norm": (["hidden_states"], ["hidden_states"]),
  56. }
  57. vocab_size: int = 50304
  58. hidden_size: int = 4096
  59. intermediate_size: int = 11008
  60. num_hidden_layers: int = 32
  61. num_attention_heads: int = 32
  62. num_key_value_heads: int | None = None
  63. hidden_act: str = "silu"
  64. max_position_embeddings: int = 2048
  65. initializer_range: float = 0.02
  66. use_cache: bool = True
  67. pad_token_id: int | None = 1
  68. bos_token_id: int | None = None
  69. eos_token_id: int | list[int] | None = 50279
  70. tie_word_embeddings: bool = False
  71. rope_parameters: RopeParameters | dict | None = None
  72. attention_bias: bool = False
  73. attention_dropout: float | int = 0.0
  74. clip_qkv: float | None = None
  75. def __post_init__(self, **kwargs):
  76. if self.num_key_value_heads is None:
  77. self.num_key_value_heads = self.num_attention_heads
  78. super().__post_init__(**kwargs)
  79. __all__ = ["OlmoConfig"]