easydel.modules.qwen2.__init__#
- class easydel.modules.qwen2.__init__.Qwen2Config(vocab_size=151936, hidden_size=4096, intermediate_size=22016, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=32, hidden_act='silu', max_position_embeddings=32768, initializer_range=0.02, rms_norm_eps=1e-06, use_cache=True, tie_word_embeddings=False, rope_theta=10000.0, use_sliding_window=False, sliding_window=4096, max_window_layers=28, attention_dropout=0.0, resid_pdrop: float = 0.0, embd_pdrop: float = 0.0, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, fcm_min_ratio: float = 0.0, fcm_max_ratio: float = 0.0, use_scan_mlp: bool = False, scan_mlp_chunk_size: int = 1024, number_rep_kv: int = 1, bits: Optional[int] = None, scan_layers: bool = True, rope_scaling: Optional[Mapping[str, str | float]] = None, **kwargs)[source]#
Bases:
EasyDeLBaseConfigConfiguration objects inherit from [PretrainedConfig] and can be used to control the model outputs. Read the documentation from [PretrainedConfig] for more information.
- Parameters
vocab_size (int, optional, defaults to 151936) – Vocabulary size of the Qwen-2 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 22016) – 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.
hidden_act (str or function, optional, defaults to “silu”) – 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 32768) – 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.
tie_word_embeddings (bool, optional, defaults to False) – 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.
use_sliding_window (bool, optional, defaults to False) – Whether to use a sliding window attention.
sliding_window (int, optional, defaults to 4096) – The sliding window size.
max_window_layers (int, optional, defaults to 28) – The maximum number of layers to use for the sliding window attention.
attention_dropout (float, optional, defaults to 0.0) – The dropout ratio for the attention probabilities.
resid_pdrop (float, optional, defaults to 0.0) – The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
embd_pdrop (float, optional, defaults to 0.0) – The dropout ratio for the embeddings.
gradient_checkpointing (str, optional, defaults to “nothing_saveable”) – The gradient checkpointing configuration.
fcm_min_ratio (float, optional, defaults to 0.0) – The minimum ratio for Flash Attention.
fcm_max_ratio (float, optional, defaults to 0.0) – The maximum ratio for Flash Attention.
use_scan_mlp (bool, optional, defaults to False) – Whether to use the scan implementation for the MLP.
scan_mlp_chunk_size (int, optional, defaults to 1024) – The chunk size to use when scanning the MLP.
number_rep_kv (int, optional, defaults to 1) – Number of repetitions for the key and value vectors.
bits (int, optional) – The number of bits to quantize the model to.
scan_layers (bool, optional, defaults to True) – Whether to use the scan implementation for the layers.
rope_scaling (tp.Dict[str, tp.Union[str, float]], optional) – The configuration for rope scaling.
- attach_custom_arguments(resid_pdrop: float = 0.0, embd_pdrop: float = 0.0, attention_dropout: float = 0.0, tie_word_embeddings: bool = False, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, fcm_min_ratio: float = 0.0, fcm_max_ratio: float = 0.0, use_scan_mlp: bool = False, scan_mlp_chunk_size: int = 1024, number_rep_kv: int = 1, bits: Optional[int] = None, rope_theta: float = 10000.0, hidden_act: str = 'silu', scan_layers: bool = True, rope_scaling: Optional[Mapping[str, str | float]] = None, **kwargs)[source]#
The attach_custom_arguments function adds the following arguments to the Transformer class:
- Parameters
self – Refer to the current object
resid_pdrop – float: Set the dropout rate for residual connections.
embd_pdrop – float: Set the probability of dropping an embedding.
attention_dropout – float: Set the probability of dropping out the attention layer.
tie_word_embeddings – bool: Tie the word embeddings to the decoder.
gradient_checkpointing – str: Control the amount of memory used by jax.
fcm_min_ratio – float: Control the minimum ratio for Flash Attention.
fcm_max_ratio – float: Set the maximum ratio for Flash Attention.
use_scan_mlp – bool: Determine whether to use the scan_mlp function or not.
scan_mlp_chunk_size – int: Set the chunk size for scan_mlp.
number_rep_kv – int: Determine how many times the key and value vectors are repeated.
bits – tp.Optional[int]: Determine the number of bits used in the quantization.
rope_theta – float: Base value for RoPE.
hidden_act – str: Activation function name.
scan_layers – bool: Determine whether to use scan layers or not.
rope_scaling (tp.Optional[tp.Mapping[str, str | float]], optional) – RoPE scaling configuration.
**kwargs – Additional keyword arguments to attach.
- get_partition_rules(fully_sharded_data_parallel: bool = True)[source]#
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()[source]#
Returns a tuple of parameter names for which weight decay should be excluded.
- property granted_freq_max_position_embedding: int#
Returns the maximum position embedding size specifically for frequency-based position embeddings.
If freq_max_position_embeddings is set, it returns that value. Otherwise, it falls back to max_position_embeddings.
- Returns
The granted maximum position embedding size for frequency encoding.
- Return type
int
- property granted_mask_max_position_embedding: int#
Returns the maximum position embedding size specifically for mask-based position embeddings.
If mask_max_position_embeddings is set, it returns that value. Otherwise, it falls back to max_position_embeddings.
- Returns
The granted maximum position embedding size for mask encoding.
- Return type
int
- model_type: str = 'qwen2'#
- class easydel.modules.qwen2.__init__.Qwen2ForCausalLM(*args: Any, **kwargs: Any)[source]#
Bases:
EasyDeLBaseModuleQwen2 model with a Causal Language Modeling head.
This model consists of the base Qwen2 transformer (Qwen2Model) followed by a linear layer (lm_head) that projects the transformer’s output hidden states to the vocabulary size, producing logits for next token prediction. Optionally, the input token embeddings can be tied to the output projection layer.
- config#
Configuration object for the model.
- Type
- dtype#
Data type for computation.
- 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
- model#
The core Qwen2 transformer model.
- Type
- lm_head#
The linear layer for projecting hidden states to vocabulary logits.
- Type
- class easydel.modules.qwen2.__init__.Qwen2ForSequenceClassification(*args: Any, **kwargs: Any)[source]#
Bases:
EasyDeLBaseModuleQwen2 model with a Sequence Classification head.
This model consists of the base Qwen2 transformer (Qwen2Model) followed by a linear layer (score) that projects the transformer’s output hidden states (typically the hidden state of the last token or a pooled representation) to the number of classes for classification.
- config#
Configuration object for the model.
- Type
- dtype#
Data type for computation.
- 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
- model#
The core Qwen2 transformer model.
- Type
- score#
The linear layer for classification.
- Type
- class easydel.modules.qwen2.__init__.Qwen2Model(*args: Any, **kwargs: Any)[source]#
Bases:
EasyDeLBaseModuleThe base Qwen2 model transformer.
This class represents the core transformer architecture of the Qwen2 model, consisting of an embedding layer, multiple Qwen2DecoderLayer layers, and a final RMS normalization layer.
- config#
Configuration object for the model.
- Type
- dtype#
Data type for computation.
- 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
- embed_tokens#
Embedding layer for input tokens.
- Type
nn.Embed
- layers#
List of decoder layers.
- Type
tp.List[Qwen2DecoderLayer]
- dropout#
Dropout layer applied after embeddings.
- Type
nn.Dropout
- gradient_checkpointing#
Gradient checkpointing configuration.