| import torch |
| from transformers import PretrainedConfig, AutoConfig |
|
|
| class MiniTransformerConfig(PretrainedConfig): |
| model_type = "minitransformer" |
|
|
| def __init__( |
| self, |
| bsz: int = 1, |
| dim: int = 768, |
| num_heads: int = 24, |
| num_layers: int = 27, |
| seq_len: int = 8192, |
| window_size: int = 8192, |
| vocab_size: int = 200064, |
| mlp_scale: int = 4, |
| bias: bool = False, |
| dropout: float = 0.0, |
| softcap: float = 50.0, |
| theta: float = 10_000.0, |
| use_alibi: bool = False, |
| torch_dtype: torch.dtype = torch.bfloat16, |
| device: torch.device = None, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| self.bsz = bsz |
| self.dim = dim |
| self.num_heads = num_heads |
| self.num_layers = num_layers |
| self.seq_len = seq_len |
| self.window_size = window_size |
| self.vocab_size = vocab_size |
| self.hidden_size = dim |
| self.mlp_scale = mlp_scale |
| self.intermediate_size = self.dim * self.mlp_scale |
| self.bias = bias |
| self.dropout = dropout |
| self.softcap = softcap |
| self.theta = theta |
| self.use_alibi = use_alibi |
| self.torch_dtype = torch_dtype |
| self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
|
