easydel.modules.gemma2.__init__

class easydel.modules.gemma2.__init__.Gemma2Config(vocab_size=256000, hidden_size=3072, intermediate_size=24576, num_hidden_layers=28, num_attention_heads=16, num_key_value_heads=16, head_dim=256, hidden_activation='gelu_pytorch_tanh', max_position_embeddings=8192, initializer_range=0.02, rms_norm_eps=1e-06, use_cache=True, pad_token_id=0, eos_token_id=1, bos_token_id=2, tie_word_embeddings=True, rope_theta=10000.0, attention_bias=False, attention_dropout=0.0, final_logit_softcapping=30.0, query_pre_attn_scalar=224, sliding_window=4096, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, bits: Optional[int] = None, scan_layers: bool = False, attn_logit_softcapping: Optional[bool] = 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 256000) – Vocabulary size of the Gemma2 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 3072) – Dimensionality of the encoder layers and the pooler layer.

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

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

• hidden_activation (str or function, optional, defaults to “gelu_pytorch_tanh”) – The non-linear activation function (function or string) to use in the encoder and pooler. If string, “gelu”, “relu”, “swish” and “gelu_new” are supported.

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

• attention_bias (bool, optional, defaults to False) – Whether to use attention bias.

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

• final_logit_softcapping (float, optional, defaults to 30.0) – The soft capping value for the final logits.

• query_pre_attn_scalar (int, optional, defaults to 224) – The scalar value for the query pre-attention layer.

• sliding_window (int, optional, defaults to 4096) – The sliding window size.

• 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. :returns: The partition rules. :rtype: tp.Tuple[tp.Tuple[str, PartitionSpec]]

static get_weight_decay_exclusions()

Returns a tuple of parameter names for which weight decay should be excluded.

Returns

An empty tuple, indicating no weight decay exclusions.

Return type

tuple

property granted_freq_max_position_embedding: int

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

Returns

The maximum position embedding size, falling back to max_position_embeddings if not explicitly set.

Return type

int

property granted_mask_max_position_embedding: int

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

Returns

The maximum position embedding size, falling back to max_position_embeddings if not explicitly set.

Return type

int

model_type: str = 'gemma2'

static rng_keys()

Returns the names of the random number generator keys used by the model.

Returns

A tuple containing “params”, “dropout”, and “fcm” as the RNG keys.

Return type

tuple

class easydel.modules.gemma2.__init__.Gemma2ForCausalLM(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

Gemma2 model with a language modeling head for causal language modeling tasks.

This model extends the base Gemma2Model by incorporating a linear language modeling head on top of the base model, designed for generative tasks and text generation.

class easydel.modules.gemma2.__init__.Gemma2ForSequenceClassification(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

class easydel.modules.gemma2.__init__.Gemma2Model(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule