File size: 8,892 Bytes
82a66d3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
# Copyright 2022 MosaicML Examples authors
# SPDX-License-Identifier: Apache-2.0

"""Forked for ReplitLM"""

"""A HuggingFace-style model configuration."""


from typing import Optional, Tuple, Union
from transformers import PretrainedConfig
class ReplitLMConfig(PretrainedConfig):
    model_type = 'replit_lm'

    def __init__(
        self,
        d_model: int = 2048,
        n_heads: int = 16,
        n_layers: int = 24,
        mlp_ratio: int = 4,
        max_seq_len: int = 2048,
        vocab_size: int = 50368,
        attn_pdrop: float = 0.0,
        resid_pdrop: float = 0.0,
        emb_pdrop: float = 0.0,
        attn_impl: str = 'triton',
        attn_qk_ln: bool = False,
        attn_clip_qkv: Optional[float] = None,
        softmax_scale: Optional[float] = None,
        prefix_lm: Optional[bool] = False,
        attn_uses_sequence_id: Optional[bool] = False,
        alibi: bool = False,
        alibi_bias_max: int = 8,
        init_device: str = 'cpu',
        logit_scale: Optional[Union[float, str]] = None,
        no_bias: bool = False,
        verbose: int = 0,
        param_init_fn: str = 'kaiming_normal_',
        init_div_is_residual: Union[int, float, str, bool] = True,
        init_std: float = 0.02,
        emb_init_std: Optional[float] = None,
        emb_init_uniform_lim: Optional[Union[Tuple[float, float],
                                             float]] = None,
        init_gain: float = 0,
        fan_mode: str = 'fan_in',
        init_nonlinearity: str = 'relu',
        embedding_fraction: float = 1.0,
        low_precision_layernorm: bool = True,
        use_cache: bool = False,
        **kwargs,
    ):
        """The ReplitLM configuration class.

        Args:
            d_model (int): The size of the embedding dimension of the model.
            n_heads (int): The number of attention heads.
            n_layers (int): The number of layers in the model.
            mlp_ratio (int): The ratio of the up/down scale in the MLP.
            max_seq_len (int): The maximum sequence length of the model.
            vocab_size (int): The size of the vocabulary.
            attn_pdrop (float): The dropout probability for the attention layers.
            resid_pdrop (float): The dropout probability applied to the attention output before combining with residual.
            emb_pdrop (float): The dropout probability for the embedding layer.
            attn_impl (str): The attention implementation to use. One of 'torch', 'flash', or 'triton'.
            attn_qk_ln (bool): Whether to apply layer normalization to the queries and keys in the attention layer.
            attn_clip_qkv (Optional[float]): If not None, clip the queries, keys, and values in the attention layer to
                this value.
            softmax_scale (Optional[float]): If not None, scale the softmax in the attention layer by this value. If None,
                use the default scale of ``1/sqrt(d_keys)``.
            prefix_lm (Optional[bool]): Whether the model should operate as a Prefix LM. This requires passing an
                extra `prefix_mask` argument which indicates which tokens belong to the prefix. Tokens in the prefix
                can attend to one another bi-directionally. Tokens outside the prefix use causal attention.
            attn_uses_sequence_id (Optional[bool]): Whether to restrict attention to tokens that have the same sequence_id.
                When the model is in `train` mode, this requires passing an extra `sequence_id` argument which indicates
                which sub-sequence each token belongs to.
                Defaults to ``False`` meaning any provided `sequence_id` will be ignored.
            alibi (bool): Whether to use the alibi bias instead of position embeddings.
            alibi_bias_max (int): The maximum value of the alibi bias.
            init_device (str): The device to use for parameter initialization.
            logit_scale (Optional[Union[float, str]]): If not None, scale the logits by this value.
            no_bias (bool): Whether to use bias in all layers.
            verbose (int): The verbosity level. 0 is silent.
            param_init_fn (str): The parameter initialization scheme to use. One of 'default_', 'baseline_', 'kaiming_uniform_',
                'kaiming_normal_', 'neox_init_', 'small_init_', 'xavier_uniform_', or 'xavier_normal_'.
            init_div_is_residual (Union[int, float, str, bool]): Value to divide initial weights by if ``module._is_residual`` is True.
            init_std (float): The standard deviation of the normal distribution used to initialize the model,
                if using the baseline_ parameter initialization scheme.
            emb_init_std (Optional[float]): The standard deviation of the normal distribution used to initialize the embedding layer.
            emb_init_uniform_lim (Optional[Union[Tuple[float, float], float]]): The lower and upper limits of the uniform distribution
                used to initialize the embedding layer. Mutually exclusive with ``emb_init_std``.
            init_gain (float): The gain to use for parameter initialization with kaiming or xavier initialization schemes.
            fan_mode (str): The fan mode to use for parameter initialization with kaiming initialization schemes.
            init_nonlinearity (str): The nonlinearity to use for parameter initialization with kaiming initialization schemes.
            embedding_fraction (float): The fraction to scale the gradients of the embedding layer by.
            low_precision_layernorm (bool): Whether to use low precision layer normalization.
            use_cache (bool): Whether or not the model should return the last key/values attentions
        """
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.mlp_ratio = mlp_ratio
        self.max_seq_len = max_seq_len
        self.vocab_size = vocab_size
        self.attn_pdrop = attn_pdrop
        self.resid_pdrop = resid_pdrop
        self.emb_pdrop = emb_pdrop
        self.attn_impl = attn_impl
        self.attn_qk_ln = attn_qk_ln
        self.attn_clip_qkv = attn_clip_qkv
        self.softmax_scale = softmax_scale
        self.prefix_lm = prefix_lm
        self.attn_uses_sequence_id = attn_uses_sequence_id
        self.alibi = alibi
        self.alibi_bias_max = alibi_bias_max
        self.init_device = init_device
        self.logit_scale = logit_scale
        self.no_bias = no_bias
        self.verbose = verbose
        self.param_init_fn = param_init_fn
        self.init_div_is_residual = init_div_is_residual
        self.init_std = init_std
        self.emb_init_std = emb_init_std
        self.emb_init_uniform_lim = emb_init_uniform_lim
        self.init_std = init_std
        self.init_gain = init_gain
        self.fan_mode = fan_mode
        self.init_nonlinearity = init_nonlinearity
        self.embedding_fraction = embedding_fraction
        self.low_precision_layernorm = low_precision_layernorm
        self.use_cache = use_cache
        if 'name' in kwargs:
            del kwargs['name']
        if 'loss_fn' in kwargs:
            del kwargs['loss_fn']
        super().__init__(**kwargs)

        self._validate_config()

    def _validate_config(self):
        if self.d_model % self.n_heads != 0:
            raise ValueError('d_model must be divisible by n_heads')
        if any(prob < 0 or prob > 1
               for prob in [self.attn_pdrop, self.resid_pdrop, self.emb_pdrop]):
            raise ValueError(
                'attn_pdrop, resid_pdrop, emb_pdrop are probabilities and must be between 0 and 1'
            )
        if self.attn_impl not in ['torch', 'flash', 'triton']:
            raise ValueError(f'Unknown attn_impl={self.attn_impl}')
        if self.prefix_lm and self.attn_impl not in ['torch', 'triton']:
            raise NotImplementedError(
                'prefix_lm only implemented with torch and triton attention.')
        if self.alibi and self.attn_impl not in ['torch', 'triton']:
            raise NotImplementedError(
                'alibi only implemented with torch and triton attention.')
        if self.attn_uses_sequence_id and self.attn_impl not in [
                'torch', 'triton'
        ]:
            raise NotImplementedError(
                'attn_uses_sequence_id only implemented with torch and triton attention.'
            )
        if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
            raise ValueError(
                'model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!'
            )
        if isinstance(self.logit_scale,
                      str) and self.logit_scale != 'inv_sqrt_d_model':
            raise ValueError(
                f"{self.logit_scale=} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'."
            )