# Copyright 2022 MosaicML Examples authors # SPDX-License-Identifier: Apache-2.0 import warnings from transformers import BertConfig as TransformersBertConfig class BertConfig(TransformersBertConfig): def __init__( self, alibi_starting_size: int = 512, normalization: str = "layernorm", attention_probs_dropout_prob: float = 0.0, head_pred_act: str = "gelu", deterministic_fa2: bool = False, allow_embedding_resizing: bool = False, **kwargs, ): """Configuration class for MosaicBert. Args: alibi_starting_size (int): Use `alibi_starting_size` to determine how large of an alibi tensor to create when initializing the model. You should be able to ignore this parameter in most cases. Defaults to 512. attention_probs_dropout_prob (float): By default, turn off attention dropout in MosaicBERT Note that the custom Triton Flash Attention with ALiBi implementation does not support droput. However, Flash Attention 2 supports ALiBi and dropout https://github.com/Dao-AILab/flash-attention embed_dropout_prob (float): Dropout probability for the embedding layer. attn_out_dropout_prob (float): Dropout probability for the attention output layer. mlp_dropout_prob (float): Dropout probability for the MLP layer. allow_embedding_resizing (bool): Embeddings will be automatically resized when they are smaller than the tokenizer vocab size. """ super().__init__(attention_probs_dropout_prob=attention_probs_dropout_prob, **kwargs) self.alibi_starting_size = alibi_starting_size self.normalization = normalization self.head_pred_act = head_pred_act self.deterministic_fa2 = deterministic_fa2 self.allow_embedding_resizing = allow_embedding_resizing class FlexBertConfig(TransformersBertConfig): model_type = "flex_bert" def __init__( self, attention_layer: str = "base", attention_probs_dropout_prob: float = 0.0, attn_out_bias: bool = False, attn_out_dropout_prob: float = 0.0, attn_qkv_bias: bool = False, bert_layer: str = "prenorm", decoder_bias: bool = True, embed_dropout_prob: float = 0.0, embed_norm: bool = True, final_norm: bool = False, embedding_layer: str = "absolute_pos", encoder_layer: str = "base", loss_function: str = "cross_entropy", loss_kwargs: dict = {}, mlp_dropout_prob: float = 0.0, mlp_in_bias: bool = False, mlp_layer: str = "mlp", mlp_out_bias: bool = False, norm_kwargs: dict = {}, normalization: str = "rmsnorm", padding: str = "unpadded", head_class_act: str = "silu", head_class_bias: bool = False, head_class_dropout: float = 0.0, head_class_norm: str = False, head_pred_act: str = "silu", head_pred_bias: bool = False, head_pred_dropout: float = 0.0, head_pred_norm: bool = True, pooling_type: str = "cls", rotary_emb_dim: int | None = None, rotary_emb_base: float = 10000.0, rotary_emb_scale_base=None, rotary_emb_interleaved: bool = False, use_fa2: bool = True, use_sdpa_attn_mask: bool = False, allow_embedding_resizing: bool = False, init_method: str = "default", init_std: float = 0.02, init_cutoff_factor: float = 2.0, init_small_embedding: bool = False, initial_attention_layer: str | None = None, initial_bert_layer: str | None = None, initial_mlp_layer: str | None = None, num_initial_layers: int = 1, skip_first_prenorm: bool = False, deterministic_fa2: bool = False, sliding_window: int = -1, global_attn_every_n_layers: int = -1, local_attn_rotary_emb_base: float = -1, local_attn_rotary_emb_dim: int | None = None, unpad_embeddings: bool = False, pad_logits: bool = False, compile_model: bool = False, masked_prediction: bool = False, **kwargs, ): """ Args: attention_layer (str): Attention layer type. attention_probs_dropout_prob (float): Dropout probability for attention probabilities. attn_out_bias (bool): use bias in attention output projection. attn_out_dropout_prob (float): Dropout probability for attention output. attn_qkv_bias (bool): use bias for query, key, value linear layer(s). bert_layer (str): BERT layer type. decoder_bias (bool): use bias in decoder linear layer. embed_dropout_prob (float): Dropout probability for embeddings. embed_norm (bool): Normalize embedding output. final_norm (bool): Add normalization after the final encoder layer and before head. embedding_layer (str): Embedding layer type. encoder_layer (str): Encoder layer type. loss_function (str): Loss function to use. loss_kwargs (dict): Keyword arguments for loss function. mlp_dropout_prob (float): Dropout probability for MLP layers. mlp_in_bias (bool): Use bias in MLP input linear layer. mlp_layer (str): MLP layer type. mlp_out_bias (bool): Use bias in MLP output linear layer. norm_kwargs (dict): Keyword arguments for normalization layers. normalization (str): Normalization type. padding (str): Unpad inputs. Best with `use_fa2=True`. head_class_act (str): Activation function for classification head. head_class_bias (bool): Use bias in classification head linear layer(s). head_class_dropout (float): Dropout probability for classification head. head_class_norm (str): Normalization type for classification head. head_pred_act (str): Activation function for prediction head. head_pred_bias (bool): Use bias in prediction head linear layer(s). head_pred_dropout (float): Dropout probability for prediction head. head_pred_norm (bool): Normalize prediction head output. pooling_type (str): Pooling type. rotary_emb_dim (int | None): Rotary embedding dimension. rotary_emb_base (float): Rotary embedding base. rotary_emb_scale_base (float): Rotary embedding scale base. rotary_emb_interleaved (bool): Use interleaved rotary embeddings. use_fa2 (bool): Use FlashAttention2. Requires flash_attn package. use_sdpa_attn_mask (bool): Pass a mask to SDPA. This will prevent SDPA from using the PyTorch FA2 kernel. allow_embedding_resizing (bool): Embeddings will be automatically resized when they are smaller than the tokenizer vocab size. init_method (str): Model layers initialization method. init_std (float): Standard deviation for initialization. Used for normal and full_megatron init. init_cutoff_factor (float): Cutoff factor for initialization. Used for normal and full_megatron init. init_small_embedding (bool): Initialize embeddings with RWKV small init. initial_attention_layer (str | None): Replace first `num_initial_layers` attention_layer instance with this layer. initial_bert_layer (str | None): Replace first `num_initial_layers` bert_layer instance with this layer. initial_mlp_layer (str | None): Replace first `num_initial_layers` mlp_layer instance with this layer. num_initial_layers (int): Number of initial layers to set via `initial_attention_layer`, `initial_bert_layer`, and `initial_mlp_layer`. skip_first_prenorm (bool): Skip pre-normalization for the first bert layer. Requires `embed_norm=True`. deterministic_fa2 (bool): Use Flash Attention 2 deterministic mode. This is slower then the default non-deterministic mode. sliding_window (int): Use sliding window attention with window size `n`. -1 to disable. Window size split between the left and right context. Only supports FA2. global_attn_every_n_layers (int): Use global attention every `n` layers and sliding window for the rest. -1 to disable. local_attn_rotary_emb_base (float): Rotary embedding base for local attention. -1 to disable and use `rotary_emb_base` for all layers. local_attn_rotary_emb_dim (int | None): Rotary embedding dimension for local attention. None to disable and use `rotary_emb_dim` for all layers. unpad_embeddings (bool): Unpad inputs before the embedding layer. pad_logits (bool): Pad logits after the calculating the loss. compile_model (bool): Compile the subset of the model which can be compiled. masked_prediction (bool): Use only pass the masked tokens throught the final MLM layers **kwargs: Additional keyword arguments. """ super().__init__(attention_probs_dropout_prob=attention_probs_dropout_prob, **kwargs) self.attention_layer = attention_layer self.attn_out_bias = attn_out_bias self.attn_out_dropout_prob = attn_out_dropout_prob self.attn_qkv_bias = attn_qkv_bias self.bert_layer = bert_layer self.decoder_bias = decoder_bias self.embed_dropout_prob = embed_dropout_prob self.embed_norm = embed_norm self.final_norm = final_norm self.embedding_layer = embedding_layer self.encoder_layer = encoder_layer self.loss_function = loss_function self.loss_kwargs = loss_kwargs self.mlp_dropout_prob = mlp_dropout_prob self.mlp_in_bias = mlp_in_bias self.mlp_layer = mlp_layer self.mlp_out_bias = mlp_out_bias self.norm_kwargs = norm_kwargs self.normalization = normalization self.padding = padding self.head_class_act = head_class_act self.head_class_bias = head_class_bias self.head_class_dropout = head_class_dropout self.head_class_norm = head_class_norm self.head_pred_act = head_pred_act self.head_pred_bias = head_pred_bias self.head_pred_dropout = head_pred_dropout self.head_pred_norm = head_pred_norm self.pooling_type = pooling_type self.rotary_emb_dim = rotary_emb_dim self.rotary_emb_base = rotary_emb_base self.rotary_emb_scale_base = rotary_emb_scale_base self.rotary_emb_interleaved = rotary_emb_interleaved self.use_fa2 = use_fa2 self.use_sdpa_attn_mask = use_sdpa_attn_mask self.allow_embedding_resizing = allow_embedding_resizing self.init_method = init_method self.init_std = init_std self.init_cutoff_factor = init_cutoff_factor self.init_small_embedding = init_small_embedding self.initial_attention_layer = initial_attention_layer self.initial_bert_layer = initial_bert_layer self.initial_mlp_layer = initial_mlp_layer self.num_initial_layers = num_initial_layers self.skip_first_prenorm = skip_first_prenorm self.deterministic_fa2 = deterministic_fa2 self.sliding_window = sliding_window self.global_attn_every_n_layers = global_attn_every_n_layers self.local_attn_rotary_emb_base = local_attn_rotary_emb_base self.local_attn_rotary_emb_dim = local_attn_rotary_emb_dim self.unpad_embeddings = unpad_embeddings self.pad_logits = pad_logits self.compile_model = compile_model self.masked_prediction = masked_prediction if loss_kwargs.get("return_z_loss", False): if loss_function != "fa_cross_entropy": raise ValueError("loss_function must be 'fa_cross_entropy' when return_z_loss is True") if loss_kwargs.get("lse_square_scale", 0) <= 0: raise ValueError( "lse_square_scale must be passed to `loss_kwargs` and must be greater than 0 for z_loss" ) if loss_kwargs.get("inplace_backward", False): self.loss_kwargs["inplace_backward"] = False warnings.warn("`inplace_backward=True` will cause incorrect metrics. Automatically setting to False.") if global_attn_every_n_layers > 0 and (self.num_hidden_layers - 1) % global_attn_every_n_layers != 0: raise ValueError( f"{global_attn_every_n_layers=} must be a divisor of one less than {self.num_hidden_layers=}" ) if self.sliding_window != -1: if not self.use_fa2: raise ValueError("Sliding window attention is only supported with FlashAttention2") if self.sliding_window % 2 != 0 and self.sliding_window % 64 != 0: raise ValueError( f"Sliding window must be an even number and divisible by 64: {self.sliding_window=} {self.sliding_window % 64} {self.sliding_window % 2}" ) else: if self.global_attn_every_n_layers != -1: raise ValueError("global_attn_every_n_layers must be -1 when sliding_window is disabled") if self.local_attn_rotary_emb_base != -1: raise ValueError("local_attn_rotary_emb_base must be -1 when sliding_window is disabled") if self.local_attn_rotary_emb_dim is not None: raise ValueError("local_attn_rotary_emb_dim must be None when sliding_window is disabled") if self.unpad_embeddings and self.padding != "unpadded": warnings.warn( "`unpad_embeddings=True` requires `padding='unpadded'`. Automatically setting `padding='unpadded'`." ) self.padding = "unpadded" if self.pad_logits and not self.unpad_embeddings: raise ValueError("`pad_logits=True` requires `unpad_embeddings=True`") if self.unpad_embeddings and self.embedding_layer == "absolute_pos": raise ValueError(f"{self.unpad_embeddings=} is incompatible with {self.embedding_layer=}") PADDING = ["unpadded", "padded"] def maybe_add_padding(config: FlexBertConfig, config_option: str) -> str: if config.padding not in PADDING: raise ValueError(f"Invalid padding type: {config.padding}, must be one of {PADDING}") if not any(config_option.startswith(pad + "_") for pad in PADDING): config_option = f"{config.padding}_{config_option}" return config_option