easydel.modules.olmo2.__init__

class easydel.modules.olmo2.__init__.Olmo2Config(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-05, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, use_scan_mlp: bool = False, scan_mlp_chunk_size: int = 1024, bits: Optional[int] = None, **kwargs)

Bases: EasyDeLBaseConfig

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. Instantiating a configuration with the defaults will yield a similar configuration to that of the [allenai/Olmo2-7B-1124-hf](https://huggingface.co/allenai/Olmo2-7B-1124-hf).

Configuration objects inherit from [PretrainedConfig] and can be used to control the model outputs. Read the documentation from [PretrainedConfig] for more information.

Parameters

• vocab_size (int, optional, defaults to 50304) – Vocabulary size of the Olmo2 model. Defines the number of different tokens that can be represented by the inputs_ids passed when calling [Olmo2Model]

• 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 in the Transformer decoder.

• num_key_value_heads (int, optional) – This is the number of key_value heads that should be used to implement Grouped Query Attention. If num_key_value_heads=num_attention_heads, the model will use Multi Head Attention (MHA), if num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed by meanpooling all the original heads within that group. For more details checkout [this paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to num_attention_heads.

• hidden_act (str or function, optional, defaults to “silu”) – The non-linear activation function (function or string) in the decoder.

• max_position_embeddings (int, optional, defaults to 2048) – The maximum sequence length that this model might ever be used with.

• initializer_range (float, optional, defaults to 0.02) – The standard deviation of the truncated_normal_initializer for initializing all weight matrices.

• use_cache (bool, optional, defaults to True) – Whether or not the model should return the last key/values attentions (not used by all models). Only relevant if config.is_decoder=True.

• 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 (tp.Dict, optional) – Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is {“type”: strategy name, “factor”: scaling factor}. When using this flag, don’t update max_position_embeddings to the expected new maximum. See the following thread for more information on how these scaling strategies behave: https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an experimental feature, subject to breaking API changes in future versions.

• attention_bias (bool, defaults to False, optional, defaults to False) – Whether to use a bias in the query, key, value and output projection layers during self-attention.

• 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.

>>> from transformers import Olmo2Model, Olmo2Config

>>> # Initializing a Olmo2 7B style configuration
>>> configuration = Olmo2Config()

>>> # Initializing a model from the Olmo2 7B style configuration
>>> model = Olmo2Model(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config

attach_custom_arguments(gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, use_scan_mlp: bool = False, scan_mlp_chunk_size: int = 1024, bits: Optional[int] = None)

get_partition_rules(*args, **kwargs)

Get the partition rules for the model. :returns: The partition rules. :rtype: tp.Tuple[tp.Tuple[str, PartitionSpec]]

property granted_freq_max_position_embedding: int

property granted_mask_max_position_embedding: int

keys_to_ignore_at_inference = ['past_key_values']

model_type: str = 'olmo2'

class easydel.modules.olmo2.__init__.Olmo2ForCausalLM(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

class easydel.modules.olmo2.__init__.Olmo2Model(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule