Repository navigation
Feature/support olmo3 #4445
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: develop
Are you sure you want to change the base?
Feature/support olmo3 #4445
Changes from 12 commits
542dee9
3c13c64
687ad58
ceff9b8
0a0c65b
dc0a377
7118fe5
9530a15
5b4d97c
2c7ac27
6f247b2
9ff6a1b
5fa6263
1dee929
7b761e1
4d97691
c555f60
a48b9e4
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,39 @@ | ||
| # Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. | ||
| # | ||
| # 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. | ||
|
|
||
| import sys | ||
| from typing import TYPE_CHECKING | ||
|
|
||
| from ...utils.lazy_import import _LazyModule | ||
|
|
||
| import_structure = { | ||
| "configuration": ["Olmo2Config"], | ||
| "modeling": [ | ||
| "Olmo2PretrainedModel", | ||
| "Olmo2Model", | ||
| "Olmo2ForCausalLM", | ||
| "Olmo2ForCausalLMPipe", | ||
| ], | ||
| } | ||
|
|
||
| if TYPE_CHECKING: | ||
| from .configuration import * | ||
| from .modeling import * | ||
| else: | ||
| sys.modules[__name__] = _LazyModule( | ||
| __name__, | ||
| globals()["__file__"], | ||
| import_structure, | ||
| module_spec=__spec__, | ||
| ) |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,149 @@ | ||
| # Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. | ||
| # Copyright 2025 the HuggingFace Team. All rights reserved. | ||
| # | ||
| # 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. | ||
|
|
||
| from ..configuration_utils import PretrainedConfig | ||
| from ..modeling_rope_utils import rope_config_validation, standardize_rope_params | ||
|
|
||
|
|
||
| class Olmo2Config(PretrainedConfig): | ||
| r""" | ||
| This is the configuration class to store the configuration of a [`Olmo2Model`]. | ||
| It is used to instantiate an OLMo2 model according to the specified arguments, | ||
| defining the model architecture. | ||
|
|
||
| Configuration objects inherit from [`PretrainedConfig`] and can be used to | ||
| control the model outputs. Read the documentation from [`PretrainedConfig`] | ||
| for more information. | ||
|
|
||
| Args: | ||
| vocab_size (`int`, *optional*, defaults to 50304): | ||
| Vocabulary size of the Olmo2 model. | ||
| hidden_size (`int`, *optional*, defaults to 4096): | ||
| Dimension of the hidden representations. | ||
| intermediate_size (`int`, *optional*, defaults to 11008): | ||
| Dimension of the MLP representations. | ||
| num_hidden_layers (`int`, *optional*, defaults to 32): | ||
| Number of hidden layers in the Transformer decoder. | ||
| num_attention_heads (`int`, *optional*, defaults to 32): | ||
| Number of attention heads for each attention layer. | ||
| num_key_value_heads (`int`, *optional*): | ||
| Number of key_value heads for Grouped Query Attention. | ||
| hidden_act (`str`, *optional*, defaults to `"silu"`): | ||
| The non-linear activation function in the decoder. | ||
| max_position_embeddings (`int`, *optional*, defaults to 2048): | ||
| The maximum sequence length. | ||
| initializer_range (`float`, *optional*, defaults to 0.02): | ||
| The standard deviation of the truncated_normal_initializer. | ||
| use_cache (`bool`, *optional*, defaults to `True`): | ||
| Whether or not the model should return last key/values attentions. | ||
| pad_token_id (`int`, *optional*, defaults to 1): | ||
| Padding token id. | ||
| bos_token_id (`int`, *optional*): | ||
| Beginning of stream token id. | ||
| eos_token_id (`int`, *optional*, defaults to 50279): | ||
| End of stream token id. | ||
| tie_word_embeddings (`bool`, *optional*, defaults to `False`): | ||
| Whether to tie weight embeddings. | ||
| rope_theta (`float`, *optional*, defaults to 10000.0): | ||
| The base period of the RoPE embeddings. | ||
| rope_scaling (`Dict`, *optional*): | ||
| Dictionary containing the scaling configuration for the RoPE embeddings. | ||
| attention_bias (`bool`, *optional*, defaults to `False`): | ||
| Whether to use a bias in the attention projection layers. | ||
| attention_dropout (`float`, *optional*, defaults to 0.0): | ||
| The dropout ratio for the attention probabilities. | ||
| rms_norm_eps (`float`, *optional*, defaults to 1e-05): | ||
| The epsilon used by the rms normalization layers. | ||
| """ | ||
|
|
||
| model_type = "olmo2" | ||
| keys_to_ignore_at_inference = ["past_key_values"] | ||
|
|
||
| def __init__( | ||
| self, | ||
| vocab_size=50304, | ||
| hidden_size=4096, | ||
| intermediate_size=11008, | ||
| num_hidden_layers=32, | ||
| num_attention_heads=32, | ||
| num_key_value_heads=None, | ||
| hidden_act="silu", | ||
| max_position_embeddings=2048, | ||
| initializer_range=0.02, | ||
| use_cache=True, | ||
| pad_token_id=1, | ||
| bos_token_id=None, | ||
| eos_token_id=50279, | ||
| tie_word_embeddings=False, | ||
| rope_theta=10000.0, | ||
| rope_scaling=None, | ||
| attention_bias=False, | ||
| attention_dropout=0.0, | ||
| rms_norm_eps=1e-5, | ||
| **kwargs, | ||
| ): | ||
| self.vocab_size = vocab_size | ||
| self.max_position_embeddings = max_position_embeddings | ||
| self.hidden_size = hidden_size | ||
| self.intermediate_size = intermediate_size | ||
| self.num_hidden_layers = num_hidden_layers | ||
| self.num_attention_heads = num_attention_heads | ||
|
|
||
| if num_key_value_heads is None: | ||
| num_key_value_heads = num_attention_heads | ||
|
|
||
| self.num_key_value_heads = num_key_value_heads | ||
| self.hidden_act = hidden_act | ||
| self.initializer_range = initializer_range | ||
| self.use_cache = use_cache | ||
| self.attention_bias = attention_bias | ||
| self.attention_dropout = attention_dropout | ||
| self.rms_norm_eps = rms_norm_eps | ||
|
|
||
| self.diff_align_float32 = kwargs.pop("diff_align_float32", False) | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
对齐实验直接用 |
||
|
|
||
| self.rope_theta = rope_theta | ||
| self.rope_scaling = rope_scaling | ||
| if self.rope_scaling is not None and "type" in self.rope_scaling: | ||
| self.rope_scaling["rope_type"] = self.rope_scaling["type"] | ||
| self.rope_parameters = self.rope_scaling | ||
| standardize_rope_params(self, rope_theta=self.rope_theta) | ||
| rope_config_validation(self) | ||
|
|
||
| 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, | ||
| ) | ||
|
|
||
| def _rope_scaling_validation(self): | ||
| """Validate the `rope_scaling` configuration.""" | ||
| if self.rope_scaling is None: | ||
| return | ||
|
|
||
| if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2: | ||
| raise ValueError( | ||
| f"`rope_scaling` must be a dictionary with two fields, `type` and `factor`, got {self.rope_scaling}" | ||
| ) | ||
| rope_scaling_type = self.rope_scaling.get("type", None) | ||
| rope_scaling_factor = self.rope_scaling.get("factor", None) | ||
| if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]: | ||
| raise ValueError( | ||
| f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}" | ||
| ) | ||
| if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0: | ||
| raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}") | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
没注册Olmo3ForCausalLM,虽然目前可以正常跑,但为了一致性需要补一下
同时也需要在docs/zh/model_capability.md和仓库README补一下该模型相关的内容