| import torch |
| from transformers import PretrainedConfig, AutoConfig |
|
|
| class MiniTransformerConfig(PretrainedConfig): |
| model_type = "minitransformer" |
|
|
| def __init__( |
| self, |
| bsz: int = 1, |
| n_embd: int = 768, |
| n_heads: int = 8, |
| n_layers: int = 25, |
| seq_len: int = 8192, |
| window_size: int = 1024, |
| vocab_size: int = 200064, |
| mlp_scale: int = 12, |
| bias: bool = False, |
| dropout: float = 0.0, |
| softcap: float = 50.0, |
| torch_dtype = torch.bfloat16, |
| device: str = None, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| self.bsz = bsz |
| self.n_embd = n_embd |
| self.n_heads = n_heads |
| self.n_layers = n_layers |
| self.seq_len = seq_len |
| self.window_size = window_size |
| self.vocab_size = vocab_size |
| self.hidden_size = n_embd |
| self.intermediate_size = n_embd * mlp_scale |
| self.hidden_act = "swish" |
| self.bias = bias |
| self.dropout = dropout |
| self.softcap = softcap |
| self.torch_dtype = torch_dtype |
| self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
