easydel.modules.gemma3.__init__

easydel.modules.gemma3.__init__#

class easydel.modules.gemma3.__init__.Gemma3Config(text_config: Optional[Gemma3TextConfig] = None, vision_config: Optional[SiglipVisionConfig] = None, mm_tokens_per_image: int = 256, boi_token_index: int = 255999, eoi_token_index: int = 256000, image_token_index: int = 262144, initializer_range: float = 0.02, **kwargs)[source]#

Bases: EasyDeLBaseConfig

Parameters

  • text_config (Union[Gemma3TextConfig, dict], optional) – The config object of the text backbone.

  • vision_config (Union[AutoConfig, dict], optional) – Custom vision config or dict.

  • mm_tokens_per_image (int, optional, defaults to 256) – The number of tokens per image embedding.

  • boi_token_index (int, optional, defaults to 255999) – The begin-of-image token index to wrap the image prompt.

  • eoi_token_index (int, optional, defaults to 256000) – The end-of-image token index to wrap the image prompt.

  • image_token_index (int, optional, defaults to 262144) – The image token index to encode the image prompt.

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

Example:

```python >>> from transformers import Gemma3ForConditionalGeneration, Gemma3Config, SiglipVisionConfig, Gemma3TextConfig

>>> # Initializing a Siglip-like vision config
>>> vision_config = SiglipVisionConfig()

>>> # Initializing a Gemma3 Text config
>>> text_config = Gemma3TextConfig()

>>> # Initializing a Gemma3 gemma-3-4b style configuration
>>> configuration = Gemma3Config(vision_config, text_config)

>>> # Initializing a model from the gemma-3-4b style configuration
>>> model = Gemma3TextConfig(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config
```
get_partition_rules(*args, **kwargs)[source]#

Get the partition rules for the model.

Returns

A tuple of tuples, where each inner tuple contains a regex pattern matching parameter names and the corresponding PartitionSpec for sharding those parameters across devices.

Return type

Tuple[Tuple[str, PartitionSpec]]

model_type: str = 'gemma3'#
sub_configs: dict[str, 'PretrainedConfig'] = {'text_config': <class 'easydel.modules.gemma3.gemma3_configuration.Gemma3TextConfig'>, 'vision_config': <class 'easydel.modules.siglip.configuration_siglip.SiglipVisionConfig'>}#
class easydel.modules.gemma3.__init__.Gemma3ForCausalLM(*args: Any, **kwargs: Any)[source]#

Bases: EasyDeLBaseModule

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

This model extends the base Gemma3TextModel by incorporating a linear language modeling head on top of the base model, designed for generative tasks and text generation. The model can optionally apply softcapping to logits based on configuration settings.

class easydel.modules.gemma3.__init__.Gemma3ForConditionalGeneration(*args: Any, **kwargs: Any)[source]#

Bases: EasyDeLBaseModule

get_image_features(pixel_values: Union[Array, ndarray, bool, number]) → Union[Array, ndarray, bool, number][source]#
init_cache(batch_size, max_length, starts=None, shardings=None, pad_token_id=None)[source]#

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

loss_type = 'ForCausalLM'#
prepare_inputs_for_generation(input_ids: Union[Array, ndarray, bool, number], max_length: int, pad_token_id: int, starts: int | None = None, pixel_values: Optional[Union[Array, ndarray, bool, number]] = None, attention_mask: Optional[Union[Array, ndarray, bool, number]] = None, token_type_ids: Optional[Union[Array, ndarray, bool, number]] = None)[source]#

Sets up the initial inputs required for starting autoregressive generation.

This function initializes the Key-Value cache (past_key_values) using init_cache, calculates the initial position_ids based on the input attention_mask (or assumes a contiguous range if no mask is provided), and prepares an extended attention_mask suitable for caching. It ensures inputs are placed on the correct devices/shards.

Parameters

  • input_ids (chex.Array) – The initial sequence of token IDs. Shape (batch_size, seq_length).

  • max_length (int) – The maximum sequence length that the KV cache should support.

  • pad_token_id (int) – The ID used for padding tokens. Used to calculate starts if not provided.

  • starts (int | None) – Optional pre-calculated starting positions (number of leading pads). If None, calculated using compute_prefill_length.

  • shardings (dict | None) – Optional sharding configuration passed to init_cache.

  • attention_mask (tp.Optional[chex.Array]) – An optional mask indicating which tokens should be attended to. Shape (batch_size, seq_length).

  • token_type_ids (tp.Optional[chex.Array]) – Optional segment IDs for models that use them.

Returns

A dictionary containing the prepared inputs, typically including:

  • ”past_key_values”: The initialized KV cache.

  • ”attention_mask”: The extended attention mask for generation.

  • ”position_ids”: The calculated initial position IDs.

  • ”token_type_ids”: (Optional) Prepared token type IDs.

This dictionary is then passed through prepare_inputs_for_call.

Return type

dict

update_inputs_for_generation(model_outputs, model_kwargs)[source]#

Updates the keyword arguments for the next generation step.

Specifically, it takes the past_key_values from the model_outputs of the current step and updates the model_kwargs with them. It also increments the position_ids by one for the next token prediction.

Parameters

  • model_outputs – The output object from the model’s forward pass in the previous step (should contain a past_key_values attribute).

  • model_kwargs (dict) – The dictionary of keyword arguments used for the model call. This dictionary will be modified in-place or a new one returned.

Returns

The updated model_kwargs dictionary ready for the next generation step.

Return type

dict

class easydel.modules.gemma3.__init__.Gemma3ForSequenceClassification(*args: Any, **kwargs: Any)[source]#

Bases: EasyDeLBaseModule

class easydel.modules.gemma3.__init__.Gemma3MultiModalProjector(*args: Any, **kwargs: Any)[source]#

Bases: Module

class easydel.modules.gemma3.__init__.Gemma3TextConfig(vocab_size=262208, hidden_size=2304, intermediate_size=9216, num_hidden_layers=26, num_attention_heads=8, num_key_value_heads=4, head_dim=256, hidden_activation='gelu_pytorch_tanh', max_position_embeddings=131072, 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=1000000.0, attention_bias=False, attention_dropout=0.0, query_pre_attn_scalar=256, sliding_window=4096, final_logit_softcapping=None, attn_logit_softcapping=None, cache_implementation='hybrid', rope_scaling=None, rope_local_base_freq=10000.0, sliding_window_pattern=6, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, bits: Optional[int] = None, scan_layers: bool = False, **kwargs)[source]#

Bases: EasyDeLBaseConfig

Configuration objects inherit from [EasyDeLBaseConfig] and can be used to control the model outputs. Read the documentation from [EasyDeLBaseConfig] for more information. :param vocab_size: Vocabulary size of the Gemma3Text model. Defines the number of different tokens that can be represented by the

inputs_ids passed when calling [Gemma3TextModel]

Parameters

  • hidden_size (int, optional, defaults to 2304) – Dimension of the hidden representations.

  • intermediate_size (int, optional, defaults to 9216) – Dimension of the MLP representations.

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

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

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

  • head_dim (int, optional, defaults to 256) – The attention head dimension.

  • hidden_activation (str or function, optional, defaults to “gelu_pytorch_tanh”) – The non-linear activation function (function or string) in the decoder. Will default to “gelu_pytorch_tanh” if not specified. “gelu_pytorch_tanh” uses an approximation of the “gelu” activation function.

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

  • rms_norm_eps (float, optional, defaults to 1e-06) – 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) – Padding token id.

  • eos_token_id (int, optional, defaults to 1) – End of stream token id.

  • bos_token_id (int, optional, defaults to 2) – Beginning of stream token id.

  • tie_word_embeddings (bool, optional, defaults to True) – Whether to tie weight embeddings

  • rope_theta (float, optional, defaults to 1000000.0) – The base period of the RoPE embeddings.

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

  • query_pre_attn_scalar (float, optional, defaults to 256) – Scaling factor used on the attention scores

  • sliding_window (int, optional, defaults to 4096) – in Gemma3Text, every other layer uses sliding window attention. This is the size of the sliding window.

  • final_logit_softcapping (float, optional) – Scaling factor when applying tanh softcapping on the logits.

  • attn_logit_softcapping (float, optional) – Scaling factor when applying tanh softcapping on the attention scores.

  • cache_implementation (str, optional, defaults to “hybrid”) – the cache type to be used with generate.

  • rope_scaling (Dict, optional) –

    Dictionary containing the scaling configuration for the RoPE embeddings used in gloabl attention. NOTE: if you apply new rope type and you expect the model to work on longer max_position_embeddings, we recommend you to update this value accordingly. Expected contents:

    rope_type (str):

    The sub-variant of RoPE to use. Can be one of [‘default’, ‘linear’, ‘dynamic’, ‘yarn’, ‘longrope’, ‘llama3’], with ‘default’ being the original RoPE implementation.

    factor (float, optional):

    Used with all rope types except ‘default’. The scaling factor to apply to the RoPE embeddings. In most scaling types, a factor of x will enable the model to handle sequences of length x * original maximum pre-trained length.

    original_max_position_embeddings (int, optional):

    Used with ‘dynamic’, ‘longrope’ and ‘llama3’. The original max position embeddings used during pretraining.

    attention_factor (float, optional):

    Used with ‘yarn’ and ‘longrope’. The scaling factor to be applied on the attention computation. If unspecified, it defaults to value recommended by the implementation, using the factor field to infer the suggested value.

    beta_fast (float, optional):

    Only used with ‘yarn’. Parameter to set the boundary for extrapolation (only) in the linear ramp function. If unspecified, it defaults to 32.

    beta_slow (float, optional):

    Only used with ‘yarn’. Parameter to set the boundary for interpolation (only) in the linear ramp function. If unspecified, it defaults to 1.

    short_factor (List[float], optional):

    Only used with ‘longrope’. The scaling factor to be applied to short contexts (< original_max_position_embeddings). Must be a list of numbers with the same length as the hidden size divided by the number of attention heads divided by 2

    long_factor (List[float], optional):

    Only used with ‘longrope’. The scaling factor to be applied to long contexts (< original_max_position_embeddings). Must be a list of numbers with the same length as the hidden size divided by the number of attention heads divided by 2

    low_freq_factor (float, optional):

    Only used with ‘llama3’. Scaling factor applied to low frequency components of the RoPE

    high_freq_factor (float, optional):

    Only used with ‘llama3’. Scaling factor applied to high frequency components of the RoPE

  • rope_local_base_freq (float, optional, defaults to 10000.0) – The base period of the RoPE embeddings for local attention.

  • sliding_window_pattern – Pattern for the sliding window attention.

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

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)[source]#

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

static get_weight_decay_exclusions()[source]#
property granted_freq_max_position_embedding: int#
property granted_mask_max_position_embedding: int#
model_type: str = 'gemma3_text'#
static rng_keys()[source]#
class easydel.modules.gemma3.__init__.Gemma3TextModel(*args: Any, **kwargs: Any)[source]#

Bases: EasyDeLBaseModule

property default_frequencies#
Contents