easydel.modules.openelm.openelm_configuration

class easydel.modules.openelm.openelm_configuration.OpenELMConfig(vocab_size: int = 32000, max_context_length: int = 2048, num_transformer_layers: int = 12, model_dim: int = 2048, head_dim: int = 128, qkv_multipliers: Union[Number, List[Number]] = 1.0, num_query_heads: Optional[int] = None, num_gqa_groups: int = 1, ffn_multipliers: Union[Number, List[Number]] = 4.0, ffn_with_glu: bool = True, ffn_dim_divisor: int = 256, activation_fn_name: str = 'swish', normalization_layer_name: str = 'rms_norm', normalize_qk_projections: bool = False, share_input_output_layers: bool = False, rope_freq_constant: int = 10000, rope_max_length: int = 4096, initializer_range: float = 0.02, use_cache: bool = True, bos_token_id: int = 1, eos_token_id: int = 2, rope_scaling: Dict[str, Union[str, float]] = None, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, use_scan_mlp: bool = False, scan_mlp_chunk_size: int = 1024, bits: Optional[int] = None, **kwargs)

Bases: EasyDeLBaseConfig

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

Parameters

• vocab_size (int, optional, defaults to 32000) – Vocabulary size of the OpenELM model. Defines the number of different tokens that can be represented by the inputs_ids passed to the forward method.

• max_context_length (int, optional, defaults to 2048) – The maximum sequence length that this model might ever be used with. Typically set this to something large just in case (e.g., 2048 or 4096).

• num_transformer_layers (int, optional, defaults to 12) – Number of hidden layers in the Transformer encoder.

• model_dim (int, optional, defaults to 2048) – Dimensionality of the encoder layers and the pooler layer.

• head_dim (int, optional, defaults to 128) – Dimensionality of the attention heads.

• qkv_multipliers (float or list of float, optional, defaults to 1.0) – The multiplier for the query, key, and value projections.

• num_query_heads (int, optional) – Number of query heads. If not provided, it will be calculated based on model_dim and head_dim.

• num_gqa_groups (int, optional, defaults to 1) – Number of GQA (Grouped Query Attention) groups.

• ffn_multipliers (float or list of float, optional, defaults to 4.0) – The multiplier for the feed-forward network.

• ffn_with_glu (bool, optional, defaults to True) – Whether to use a gated linear unit (GLU) in the feed-forward network.

• ffn_dim_divisor (int, optional, defaults to 256) – The divisor for the feed-forward network dimension.

• activation_fn_name (str, optional, defaults to “swish”) – The activation function to use.

• normalization_layer_name (str, optional, defaults to “rms_norm”) – The normalization layer to use.

• normalize_qk_projections (bool, optional, defaults to False) – Whether to normalize the query and key projections.

• share_input_output_layers (bool, optional, defaults to False) – Whether to share the input and output layers.

• rope_freq_constant (int, optional, defaults to 10000) – The frequency constant for Rotary Position Embeddings (RoPE).

• rope_max_length (int, optional, defaults to 4096) – The maximum length for RoPE.

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

• bos_token_id (int, optional, defaults to 1) – The id of the beginning-of-sequence token.

• eos_token_id (int, optional, defaults to 2) – The id of the end-of-sequence token.

• rope_scaling (tp.Dict[str, tp.Union[str, float]], optional) – The configuration for rope scaling.

• gradient_checkpointing (str, optional, defaults to “nothing_saveable”) – The gradient checkpointing configuration.

• use_scan_mlp (bool, optional, defaults to False) – Whether to use the scan implementation for the MLP.

• scan_mlp_chunk_size (int, optional, defaults to 1024) – The chunk size to use when scanning the MLP.

• bits (int, optional) – The number of bits to quantize the model to.

add_jax_args(gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, use_scan_mlp: bool = False, scan_mlp_chunk_size: int = 1024, bits: Optional[int] = None, rope_scaling: Dict[str, Union[str, float]] = None, **kwargs)

The add_jax_args function adds the following arguments to the model:

Parameters

• self – Bind the attributes and methods of a class to an instance of that class

• gradient_checkpointing – str: Determine whether to use gradient checkpointing

• use_scan_mlp – bool: Determine whether to use the scan_mlp function or notn

• scan_mlp_chunk_size – int: Chunk the input to the mlp

• bits – tp.Optional[int]: Specify the number of bits to use for quantization

• rope_scaling – tp.Dict[str, tp.Union[str, float]]: rope_scaling for rope

Return type

A tuple of the following

attribute_map: Dict[str, str] = {'tie_word_embedding': 'share_input_output_layers'}

get_partition_rules(*args, **kwargs)

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

static get_weight_decay_exclusions()

property granted_freq_max_position_embedding: int

property granted_mask_max_position_embedding: int

model_type: str = 'openelm'

static rng_keys()

easydel.modules.openelm.openelm_configuration.compute_heads(model_dim: int, head_dim: int) → int

Compute the number of heads. :param model_dim: Model dimension. :param head_dim: Head dimension.

Returns

An integer denoting number of heads in multi-head attention is returned.

Raises

ValueError – if model dimension is not divisible by head dimension.

easydel.modules.openelm.openelm_configuration.make_divisible(v: Union[float, int], divisor: Optional[int] = 8, min_value: Optional[Union[float, int]] = None) → Union[float, int]

This function is taken from the original tf repo. It ensures that all layers have a channel number that is divisible by the divisor It can be seen at: tensorflow/models :param v: input value :param divisor: default to 8 :param min_value: minimum divisor value

Returns

new divisible value

Return type

new_v