easydel.modules.xerxes2.__init__

class easydel.modules.xerxes2.__init__.Xerxes2Config(vocab_size: int = 256128, hidden_size: int = 4096, intermediate_size: int = 16384, num_hidden_layers: int = 32, num_attention_heads: int = 32, max_position_embeddings: int = 16384, initializer_range: float = 0.02, rms_norm_eps: float = 1e-06, use_cache: bool = True, pad_token_id: int = 0, eos_token_id: int = 1, bos_token_id: int = 2, tie_word_embeddings: bool = False, rope_theta: float = 10000.0, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, bits: Optional[int] = None, scan_layers: bool = False, q_lora_dim: Optional[int] = 1536, kv_lora_dim: int = 512, qk_rope_head_dim: int = 64, qk_nope_head_dim: int = 128, vhead_dim: int = 128, **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 256128) – Vocabulary size of the xerxes model. Defines the number of different tokens that can be represented by the inputs_ids passed to the forward method.

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

• intermediate_size (int, optional, defaults to 16384) – Dimensionality of the “intermediate” (i.e., feed-forward) layer in the Transformer encoder.

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

• num_attention_heads (int, optional, defaults to 16) – Number of attention heads for each attention layer in the Transformer encoder.

• num_key_value_heads (int, optional, defaults to 16) – Number of key and value heads for each attention layer in the Transformer encoder.

• head_dim (int, optional, defaults to 256) – Dimensionality of the attention head.

• max_position_embeddings (int, optional, defaults to 6144) – 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).

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

• rms_norm_eps (float, optional, defaults to 1e-6) – The epsilon used by the rms normalization layers.

• 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 0) – The index of the padding token in the vocabulary.

• eos_token_id (int, optional, defaults to 1) – The index of the end of sequence token in the vocabulary.

• bos_token_id (int, optional, defaults to 2) – The index of the beginning of sequence token in the vocabulary.

• tie_word_embeddings (bool, optional, defaults to True) – Whether to tie the weights of the input embeddings and the output embeddings.

• rope_theta (float, optional, defaults to 10000.0) – The theta value to use for rotary position embeddings.

• softmax_scale (float, optional, defaults to 14.9666295471) – softmax scale for attention module.

• attention_dropout (float, optional, defaults to 0.0) – The dropout ratio for the attention probabilities.

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

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

• scan_layers (bool, optional, defaults to False) – Whether to use the scan implementation of the layers.

attach_custom_arguments(gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, bits: Optional[int] = None, **kwargs)

The attach_custom_arguments function adds the following arguments to the Transformer class:

Parameters

• self – Refer to the current object

• gradient_checkpointing – str: Control the amount of memory used by jax

• bits – tp.Optional[int]: Determine the number of bits used in the quantization

get_partition_rules(*args, **kwargs)

Get the partition rules for the model.

Parameters

fully_sharded_data_parallel (bool, optional, defaults to True) – Whether to use fully sharded data parallelism.

Returns

The partition rules.

Return type

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 = 'xerxes2'

static rng_keys()

class easydel.modules.xerxes2.__init__.Xerxes2ForCausalLM(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

create_cache_metadata(batch_size: int, max_length: int, pad_token_id: int | None = None)

Creates the metadata required for initializing a standard (non-paged) KV Cache.

This method gathers parameters like layer count, head dimensions, and determines the appropriate padding token ID to instantiate and return a TransformerCacheMetaData object suitable for a standard sequential KV cache.

Parameters

• batch_size (int) – The batch size for which the cache is being configured.

• max_length (int) – The maximum sequence length the cache needs to support.

• pad_token_id (int | None) – The ID of the padding token. If None, it attempts to find it from self.generation_config or self.config, defaulting to 0.

Returns

An initialized metadata object for a standard KV cache.

Return type

TransformerCacheMetaData

init_cache(batch_size: int, max_length: int, starts: int | None = None, shardings: dict | None = None, pad_token_id: int | None = None)

Initializes and returns a standard (non-paged) Key-Value cache.

This method first creates the necessary metadata using create_cache_metadata and then calls TransformerCache.init_cache to allocate and initialize the cache tensors based on the model’s configuration, dtype, sharding, quantization settings, and provided batch size and maximum length.

Parameters

• batch_size (int) – The batch size for the cache.

• max_length (int) – The maximum sequence length the cache needs to support.

• starts (int | None) – Optional starting positions for the cache sequences. If provided, influences the initial state. Defaults to None (usually 0).

• shardings (dict | None) – Optional dictionary specifying sharding configurations. (Note: This argument appears unused in the current implementation shown).

• pad_token_id (int | None) – The ID of the padding token. If None, it’s inferred.

Returns

An initialized standard TransformerCache object.

Return type

TransformerCache

class easydel.modules.xerxes2.__init__.Xerxes2Model(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

property frequencies: Array

Returns frequency values from the config.