| 13 | |
| 14 | |
| 15 | class DecoderOnly(torch.nn.Module): |
| 16 | |
| 17 | def __init__( |
| 18 | self, |
| 19 | n_kv_head: int, |
| 20 | head_dim: int, |
| 21 | hidden_size: int, |
| 22 | attention_heads: int = 4, |
| 23 | linear_units: int = 2048, |
| 24 | num_blocks: int = 6, |
| 25 | dropout_rate: float = 0.1, |
| 26 | positional_dropout_rate: float = 0.1, |
| 27 | attention_dropout_rate: float = 0.0, |
| 28 | normalize_before: bool = True, |
| 29 | query_bias: bool = False, |
| 30 | key_bias: bool = False, |
| 31 | value_bias: bool = False, |
| 32 | mlp_bias: bool = False, |
| 33 | activation_type: str = "gelu", |
| 34 | gelu_approximate: Union[str, None] = None, |
| 35 | max_position_embeding: int = 8192, |
| 36 | mlp_type: str = 'gated', |
| 37 | layer_norm_type: str = 'rms_norm', |
| 38 | norm_eps: float = 1e-5, |
| 39 | rms_norm_offset: bool = True, |
| 40 | selfattention_layer_type: str = "rope_abs_selfattn", |
| 41 | use_sdpa: bool = False, |
| 42 | gradient_checkpointing: bool = False, |
| 43 | rope_theta: float = 10000.0, |
| 44 | rope_style: str = 'google', |
| 45 | scale_embed: bool = True, |
| 46 | ) -> None: |
| 47 | super().__init__() |
| 48 | |
| 49 | assert selfattention_layer_type in ['rope_abs_selfattn'] |
| 50 | self.pos_enc = WENET_EMB_CLASSES["rope_pos"]( |
| 51 | hidden_size, |
| 52 | head_dim, |
| 53 | max_len=max_position_embeding, |
| 54 | dropout_rate=positional_dropout_rate, |
| 55 | rope_theta=rope_theta, |
| 56 | scale=scale_embed) |
| 57 | if activation_type == "gelu" and gelu_approximate is not None: |
| 58 | activation = WENET_ACTIVATION_CLASSES['gelu']( |
| 59 | approximate=gelu_approximate) |
| 60 | else: |
| 61 | activation = WENET_ACTIVATION_CLASSES[activation_type]() |
| 62 | |
| 63 | mlp_class = WENET_MLP_CLASSES[mlp_type] |
| 64 | self.num_blocks = num_blocks |
| 65 | # TODO: support lora & refactor lora |
| 66 | self.decoders = torch.nn.ModuleList([ |
| 67 | TransformerEncoderLayer( |
| 68 | hidden_size, |
| 69 | WENET_ATTENTION_CLASSES[selfattention_layer_type]( |
| 70 | attention_heads, |
| 71 | hidden_size, |
| 72 | attention_dropout_rate, |