| 230 | """ |
| 231 | |
| 232 | def __init__( |
| 233 | self, hidden_size, compute_type=mstype.float16, parallel_config=None |
| 234 | ): |
| 235 | super(PanGuHead, self).__init__() |
| 236 | if parallel_config.vocab_emb_dp: |
| 237 | self.matmul = P.MatMul(transpose_b=True).shard( |
| 238 | ((parallel_config.data_parallel, 1), (1, 1)) |
| 239 | ) |
| 240 | else: |
| 241 | print("====vocab_emb_dp is false", flush=True) |
| 242 | self.matmul = P.MatMul(transpose_b=True).shard( |
| 243 | ( |
| 244 | (parallel_config.data_parallel, 1), |
| 245 | (parallel_config.model_parallel, 1), |
| 246 | ) |
| 247 | ) |
| 248 | self.hidden_size = hidden_size |
| 249 | self.dtype = compute_type |
| 250 | self.cast = P.Cast() |
| 251 | |
| 252 | def construct(self, state, embed): |
| 253 | state = P.Reshape()(state, (-1, self.hidden_size)) |