easydel.modules.gpt_neox.modeling_gpt_neox

easydel.modules.gpt_neox.modeling_gpt_neox#

class easydel.modules.gpt_neox.modeling_gpt_neox.GPTNeoXAttention(*args: Any, **kwargs: Any)[source]#

Bases: UnifiedAttention

GPT-NeoX Attention with partial RoPE.

Inherits from UnifiedAttention. Uses combined QKV projection (query_key_value) and partial rotary embeddings. Overrides forward_standard to efficiently handle fused QKV projection.

projection_mapping: ClassVar[dict[str, str]] = {'key_projection': 'k_proj', 'mla_kv_a_layernorm': 'kv_a_layernorm', 'mla_kv_a_proj_with_mqa': 'kv_a_proj_with_mqa', 'mla_kv_b_proj': 'kv_b_proj', 'mla_q_a_layernorm': 'q_a_layernorm', 'mla_q_a_proj': 'q_a_proj', 'mla_q_b_proj': 'q_b_proj', 'mla_q_proj': 'q_proj', 'output_projection': 'dense', 'query_key_value_projection': 'query_key_value', 'query_projection': 'q_proj', 'value_projection': 'v_proj'}#

class easydel.modules.gpt_neox.modeling_gpt_neox.GPTNeoXBlock(*args: Any, **kwargs: Any)[source]#

Bases: Module

GPT-NeoX Transformer block.

This module represents a single transformer block in the GPT-NeoX model, containing self-attention and MLP sub-layers with residual connections and layer normalization. It supports both standard and parallel residual connections.

config#

Configuration object for the model.

Type: GPTNeoXConfig

dtype#

Data type for computations.

Type: jnp.dtype

param_dtype#

Data type for parameters.

Type: jnp.dtype

precision#

Precision setting for JAX operations.

Type: jax.lax.PrecisionLike

rngs#

Random number generators.

Type: nn.Rngs

class easydel.modules.gpt_neox.modeling_gpt_neox.GPTNeoXForCausalLM(*args: Any, **kwargs: Any)[source]#

Bases: BaseCausalLMModule[GPTNeoXModel, GPTNeoXConfig]

GPT-NeoX model with a language modeling head.

class easydel.modules.gpt_neox.modeling_gpt_neox.GPTNeoXMlp(*args: Any, **kwargs: Any)[source]#

Bases: Module

GPT-NeoX MLP module.

This module implements the feed-forward network used in the GPT-NeoX model.

config#

Configuration object for the model.

Type: GPTNeoXConfig

dtype#

Data type for computations.

Type: jnp.dtype

param_dtype#

Data type for parameters.

Type: jnp.dtype

precision#

Precision setting for JAX operations.

Type: jax.lax.PrecisionLike

rngs#

Random number generators.

Type: nn.Rngs

class easydel.modules.gpt_neox.modeling_gpt_neox.GPTNeoXModel(*args: Any, **kwargs: Any)[source]#

Bases: EasyDeLBaseModule

GPT-NeoX model implementation.

This class implements the main GPT-NeoX transformer model architecture, consisting of an embedding layer, multiple GPTNeoXBlock layers, and a final layer normalization.

config#

Configuration object for the model.

Type: GPTNeoXConfig

dtype#

Data type for computations.

Type: jnp.dtype

param_dtype#

Data type for parameters.

Type: jnp.dtype

precision#

Precision setting for JAX operations.

Type: jax.lax.PrecisionLike

rngs#

Random number generators.

Type: nn.Rngs

property frequencies#

Retrieves or computes the frequency components (e.g., for RoPE) from the configuration.

Uses self.config.get_basic_frequencies() and caches the result.

Returns: The frequency components, potentially cached.
Return type: jnp.ndarray

get_decoder()[source]#: Returns the decoder part of the model’s graph definition.

get_embedding()[source]#: Returns the embedding layer of the module.

get_encoder()[source]#: Returns the encoder part of the model’s graph definition. Decoder-Only models don’t have an encoder.

get_lm_head()[source]#: Returns the language model head of the module. Base Models don’t have a Language Model Head.

easydel.modules.gpt_neox.modeling_gpt_neox

Contents

easydel.modules.gpt_neox.modeling_gpt_neox#