easydel.modules.grok_1.__init__

easydel.modules.grok_1.__init__#

class easydel.modules.grok_1.__init__.Grok1Config(vocab_size=32000, hidden_size=4096, intermediate_size=32768, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=32, attn_output_multiplier=1.0, max_attn_value=1.0, max_position_embeddings=4096, embedding_multiplier_scale: float = 1.0, output_multiplier_scale: float = 1.0, rms_norm_eps=1e-05, use_cache=True, pad_token_id=None, bos_token_id=1, eos_token_id=2, tie_word_embeddings=True, num_experts_per_tok=2, num_experts=8, output_router_logits=False, router_aux_loss_coef=0.001, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, bits: Optional[int] = None, **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.

Parameters

  • vocab_size (int, optional, defaults to 32000) – Vocabulary size of the Grok-1 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 32768) – 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 32) – Number of attention heads for each attention layer in the Transformer encoder.

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

  • attn_output_multiplier (float, optional, defaults to 1.0) – The multiplier value applied to the attention output.

  • max_attn_value (float, optional, defaults to 1.0) – The maximum value of the attention weights.

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

  • embedding_multiplier_scale (float, optional, defaults to 1.0) – The scale factor for the embedding layer.

  • output_multiplier_scale (float, optional, defaults to 1.0) – The scale factor for the output layer.

  • rms_norm_eps (float, optional, defaults to 1e-5) – 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) – The index of the padding token in the vocabulary.

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

  • eos_token_id (int, optional, defaults to 2) – The index of the end 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.

  • num_experts_per_tok (int, optional, defaults to 2) – The number of experts per token.

  • num_experts (int, optional, defaults to 8) – The number of experts.

  • output_router_logits (bool, optional, defaults to False) – Whether to output router logits.

  • router_aux_loss_coef (float, optional, defaults to 0.001) – The router auxiliary loss coefficient.

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

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

attach_custom_arguments(tie_word_embeddings: bool = False, 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

  • tie_word_embeddings – bool: Tie the word embeddings to the decoder

  • 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 = 'grok-1'#
static rng_keys()[source]#
class easydel.modules.grok_1.__init__.Grok1ForCausalLM(*args: Any, **kwargs: Any)[source]#

Bases: EasyDeLBaseModule

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

Bases: EasyDeLBaseModule

property frequencies#

Returns frequency values from the config.

Contents