configuration_olmo3.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101
  1. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
  2. # This file was automatically generated from src/transformers/models/olmo3/modular_olmo3.py.
  3. # Do NOT edit this file manually as any edits will be overwritten by the generation of
  4. # the file from the modular. If any change should be done, please apply the change to the
  5. # modular_olmo3.py file directly. One of our CI enforces this.
  6. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
  7. # Copyright 2025 the HuggingFace Team. All rights reserved.
  8. #
  9. # Licensed under the Apache License, Version 2.0 (the "License");
  10. # you may not use this file except in compliance with the License.
  11. # You may obtain a copy of the License at
  12. #
  13. # http://www.apache.org/licenses/LICENSE-2.0
  14. #
  15. # Unless required by applicable law or agreed to in writing, software
  16. # distributed under the License is distributed on an "AS IS" BASIS,
  17. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  18. # See the License for the specific language governing permissions and
  19. # limitations under the License.
  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-3-7B-Instruct")
  25. @strict
  26. class Olmo3Config(PreTrainedConfig):
  27. r"""
  28. Example:
  29. ```python
  30. >>> from transformers import Olmo3Model, Olmo3Config
  31. >>> # Initializing a Olmo3 7B style configuration
  32. >>> configuration = Olmo3Config()
  33. >>> # Initializing a model from the Olmo3 7B style configuration
  34. >>> model = Olmo3Model(configuration)
  35. >>> # Accessing the model configuration
  36. >>> configuration = model.config
  37. ```
  38. """
  39. model_type = "olmo3"
  40. keys_to_ignore_at_inference = ["past_key_values"]
  41. base_model_tp_plan = {
  42. "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k
  43. "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k
  44. "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k
  45. "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k
  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 = 50304
  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 = 0.02
  64. use_cache: bool = True
  65. pad_token_id: int | None = 1
  66. bos_token_id: int | None = None
  67. eos_token_id: int | list[int] | None = 50279
  68. tie_word_embeddings: bool = False
  69. rope_parameters: RopeParameters | dict | None = None
  70. attention_bias: bool = False
  71. attention_dropout: float | int = 0.0
  72. rms_norm_eps: float = 1e-5
  73. sliding_window: int | None = 4096
  74. layer_types: list[str] | 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. if self.layer_types is None:
  79. self.layer_types = [
  80. "sliding_attention" if (i + 1) % 4 != 0 else "full_attention" for i in range(self.num_hidden_layers)
  81. ]
  82. if self.num_key_value_heads is None:
  83. self.num_key_value_heads = self.num_attention_heads
  84. super().__post_init__(**kwargs)
  85. __all__ = ["Olmo3Config"]