easydel.modules.arctic.arctic_configuration

easydel.modules.arctic.arctic_configuration#

class easydel.modules.arctic.arctic_configuration.ArcticConfig(vocab_size=32000, hidden_size=4096, intermediate_size=14336, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=None, hidden_act='silu', max_position_embeddings=4096, initializer_range=0.02, rms_norm_eps=1e-05, use_cache=True, pad_token_id=None, bos_token_id=1, eos_token_id=2, tie_word_embeddings=False, rope_theta=1000000.0, sliding_window=None, attention_dropout=0.0, num_experts_per_tok=1, num_local_experts=8, router_aux_loss_coef=0.001, moe_layer_frequency=2, parallel_attn_mlp_res=False, moe_train_capacity_factor=1, moe_eval_capacity_factor=1, enable_expert_tensor_parallelism=False, moe_min_capacity=0, moe_token_dropping=True, quantization=None, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, use_scan_mlp: bool = False, scan_mlp_chunk_size: int = 1024, layer_types: list[str] | None = None, bits: int | None = None, rope_scaling: dict[str, str | float] | None = 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 ARCTIC 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 14336) – 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) – Number of key and value heads for each attention layer in the Transformer encoder. Will default to num_attention_heads if not set.

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

  • 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-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. The default value (0) is the same as for GPT2.

  • bos_token_id (int, optional) – The index of the beginning of sequence token in the vocabulary. The default value (1) is the same as for GPT2.

  • eos_token_id (int, optional) – The index of the end of sequence token in the vocabulary. The default value (2) is the same as for GPT2.

  • 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 1e6) – The theta value to use for rotary position embeddings.

  • sliding_window (int, optional) – The sliding window size to use for attention. If not specified, no sliding window attention is used.

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

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

  • num_local_experts (int, optional, defaults to 8) – The number of local experts for mixture of experts.

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

  • moe_layer_frequency (int, optional, defaults to 2) – The frequency of MoE layers.

  • parallel_attn_mlp_res (bool, optional, defaults to False) – Whether to parallelize attention and MLP residual connections.

  • moe_train_capacity_factor (float, optional, defaults to 1) – The capacity factor for MoE layers during training.

  • moe_eval_capacity_factor (float, optional, defaults to 1) – The capacity factor for MoE layers during evaluation.

  • enable_expert_tensor_parallelism (bool, optional, defaults to False) – Whether to enable expert tensor parallelism.

  • moe_min_capacity (int, optional, defaults to 0) – The minimum capacity for MoE layers.

  • moe_token_dropping (bool, optional, defaults to True) – Whether to drop tokens in MoE layers.

  • quantization (str, optional) – The quantization configuration.

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

  • use_scan_mlp (bool, optional, defaults to False) – Whether to use scan for MLP.

  • scan_mlp_chunk_size (int, optional, defaults to 1024) – The chunk size for scan MLP.

  • bits (int, optional) – The number of bits.

  • rope_scaling (tp.Dict[str, tp.Union[str, float]], optional) – The rope scaling configuration.

get_mask_details() → dict[int, easydel.infra.utils.AttnMaskDetail][source]#

Retrieve attention mask details for each layer in the model.

This method generates a dictionary mapping layer indices to their corresponding attention mask details. If a sliding window is defined, each layer is assigned a sliding window attention mask with the specified size.

Returns

A dictionary where keys are layer indices (int) and values are AttnMaskDetail objects specifying the attention mask type and size for each layer.

Return type

dict[int, AttnMaskDetail]

Notes

  • If self.sliding_window is None, an empty dictionary is returned.

  • The method iterates over self.num_hidden_layers to assign mask details for each layer.

  • The attention mask type is set to AttnMaskType.SLIDING when a sliding window is defined.

get_partition_rules(*args, **kwargs)[source]#

Shard parameters for Arctic model with MoE support.

model_type: str = 'arctic'#
class easydel.modules.arctic.arctic_configuration.ExpertTensorParallel(axes: Sequence[str | None], mode: Union[Literal['__autoregressive__', '__prefill__', '__train__', '__insert__'], int])[source]#

Bases: DynamicShardingAxes

Expert Tensor Parallelism (EPxTP) sharding axes.

axes: ClassVar = ['__TENSOR_PARALLEL__', '_', '_']#
mode: ClassVar = '__train__'#
Contents