| 176 | |
| 177 | |
| 178 | class MossMLP(Module): |
| 179 | def __init__(self, intermediate_size, config): |
| 180 | # in MLP: intermediate_size= 4 * embed_dim |
| 181 | super(MossMLP, self).__init__() |
| 182 | embed_dim = config.n_embd |
| 183 | |
| 184 | self.fc_in = nn.Linear(embed_dim, intermediate_size) |
| 185 | self.fc_out = nn.Linear(intermediate_size, embed_dim) |
| 186 | |
| 187 | self.act = NewGELUActivation() |
| 188 | self.dropout = nn.Dropout(config.resid_pdrop) |
| 189 | |
| 190 | def execute(self, hidden_states: Optional[jt.Var]) -> jt.Var: |
| 191 | hidden_states = self.fc_in(hidden_states) |
| 192 | hidden_states = self.act(hidden_states) |
| 193 | hidden_states = self.fc_out(hidden_states) |
| 194 | hidden_states = self.dropout(hidden_states) |
| 195 | return hidden_states |
| 196 | |
| 197 | |
| 198 | class MossBlock(Module): |