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)

Attaches custom arguments to the configuration object.

This method allows adding or overriding configuration attributes dynamically. It primarily sets attributes related to gradient checkpointing, MLP scanning, and quantization bits.

Parameters

• gradient_checkpointing (EasyDeLGradientCheckPointers, optional) – Gradient checkpointing strategy. Defaults to EasyDeLGradientCheckPointers.NONE.

• use_scan_mlp (bool, optional) – Whether to use scan for MLP layers. Defaults to False.

• scan_mlp_chunk_size (int, optional) – Chunk size for scan MLP. Defaults to 1024.

• bits (tp.Optional[int], optional) – Quantization bits. Defaults to None.

get_partition_rules(*args, **kwargs)

Get the partition rules for the model. This method defines how the model’s parameters are partitioned across devices for distributed training and inference.

Parameters

• *args – Additional positional arguments (unused).

• **kwargs – Additional keyword arguments (unused).

Returns

A tuple of partition rules, where each rule is a tuple

containing a regex pattern for parameter names and the corresponding PartitionSpec.

Return type

tp.Tuple[tp.Tuple[str, PartitionSpec]]

property granted_freq_max_position_embedding: int

Returns the maximum position embedding size specifically for frequency-based position embeddings.

If freq_max_position_embeddings is set, it returns that value. Otherwise, it falls back to max_position_embeddings.

Returns

The granted maximum position embedding size for frequency encoding.

Return type

int

property granted_mask_max_position_embedding: int

Returns the maximum position embedding size specifically for mask-based position embeddings.

If mask_max_position_embeddings is set, it returns that value. Otherwise, it falls back to max_position_embeddings.

Returns

The granted maximum position embedding size for mask encoding.

Return type

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

OLMo-2 model with a Causal Language Modeling head.

This model consists of the base OLMo-2 transformer (Olmo2Model) followed by a linear layer (lm_head) that projects the transformer’s output hidden states to the vocabulary size, producing logits for next token prediction.

config

Configuration object for the model.

Type

Olmo2Config

dtype

Data type for computation.

Type

jnp.dtype

param_dtype

Data type for parameters.

Type

jnp.dtype

precision

Precision setting for JAX operations.

Type

jax.lax.PrecisionLike

rngs

Random number generators.

Type

nn.Rngs

model

The core OLMo-2 transformer model.

Type

Olmo2Model

lm_head

The linear layer for projecting hidden states to vocabulary logits.

Type

ParallelLinear

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

Bases: EasyDeLBaseModule

The base OLMo-2 model transformer.

This class represents the core transformer architecture of the OLMo-2 model, consisting of an embedding layer, multiple Olmo2DecoderLayer layers, and a final RMS normalization layer.

config

Configuration object for the model.

Type

Olmo2Config

dtype

Data type for computation.

Type

jnp.dtype

param_dtype

Data type for parameters.

Type

jnp.dtype

precision

Precision setting for JAX operations.

Type

jax.lax.PrecisionLike

rngs

Random number generators.

Type

nn.Rngs

embed_tokens

Embedding layer for input tokens.

Type

nn.Embed

layers

List of decoder layers.

Type

tp.List[Olmo2DecoderLayer]

norm

Final layer normalization.

Type

RMSNorm

gradient_checkpointing

Gradient checkpointing configuration.

Type

EasyDeLGradientCheckPointers