| 191 | |
| 192 | |
| 193 | class T5Model(EncoderDecoderModel): |
| 194 | def __init__(self, args, **kwargs): |
| 195 | self.init_method_std = args.init_method_std |
| 196 | super().__init__(args, tie_word_embeddings=True, **kwargs, use_bias=False, |
| 197 | layernorm=T5LayerNorm, activation_func=torch.nn.functional.relu, |
| 198 | init_method=self._init_weights) |
| 199 | self.encoder.add_mixin( |
| 200 | "t5-attention", T5AttentionMixin(args.relative_attention_num_buckets, args.num_attention_heads) |
| 201 | ) |
| 202 | self.encoder.add_mixin( |
| 203 | "t5-position", T5PositionEmbeddingMixin() |
| 204 | ) |
| 205 | del self.encoder.transformer.position_embeddings |
| 206 | num_attention_heads = args.dec_num_attention_heads if args.dec_num_attention_heads is not None else args.num_attention_heads |
| 207 | self.decoder.add_mixin( |
| 208 | "t5-attention", T5AttentionMixin(args.relative_attention_num_buckets, num_attention_heads, is_decoder=True) |
| 209 | ) |
| 210 | self.decoder.add_mixin( |
| 211 | "t5-position", T5PositionEmbeddingMixin() |
| 212 | ) |
| 213 | self.decoder.add_mixin( |
| 214 | "t5-final", |
| 215 | T5DecoderFinalMixin(args.vocab_size, args.hidden_size, tie_word_embeddings=not args.no_share_embeddings) |
| 216 | ) |
| 217 | del self.decoder.transformer.position_embeddings |
| 218 | if args.gated_gelu_mlp: |
| 219 | self.encoder.add_mixin( |
| 220 | "gated-mlp", T5GatedGeluMLPMixin(args.num_layers, args.hidden_size, init_method_std=self.init_method_std, |
| 221 | inner_hidden_size=args.inner_hidden_size, bias=False) |
| 222 | ) |
| 223 | self.decoder.add_mixin( |
| 224 | "gated-mlp", T5GatedGeluMLPMixin(args.num_layers, args.hidden_size, init_method_std=self.init_method_std, |
| 225 | inner_hidden_size=args.inner_hidden_size, bias=False) |
| 226 | ) |
| 227 | |
| 228 | def _init_weights(self, weight, module, name): |
| 229 | init_method_std = self.init_method_std |
| 230 | if isinstance(module, MLP): |
| 231 | if name == "dense_h_to_4h": |
| 232 | torch.nn.init.normal_(weight, mean=0, std=init_method_std * (module.hidden_size ** -0.5)) |
| 233 | elif name == "dense_4h_to_h": |
| 234 | torch.nn.init.normal_(weight, mean=0, std=init_method_std * (module.inner_hidden_size ** -0.5)) |
| 235 | else: |
| 236 | raise NotImplementedError(name) |
| 237 | elif isinstance(module, SelfAttention): |
| 238 | if name == "query_key_value": |
| 239 | torch.nn.init.normal_(weight, mean=0, std=init_method_std * (module.hidden_size ** -0.5)) |
| 240 | torch.nn.init.normal_(weight[:module.inner_hidden_size], mean=0, std=init_method_std * ( |
| 241 | (module.hidden_size * module.hidden_size_per_attention_head) ** -0.5)) |
| 242 | elif name == "dense": |
| 243 | torch.nn.init.normal_(weight, mean=0, std=init_method_std * (module.inner_hidden_size ** -0.5)) |
| 244 | else: |
| 245 | raise NotImplementedError(name) |
| 246 | elif isinstance(module, CrossAttention): |
| 247 | if name == "query": |
| 248 | torch.nn.init.normal_(weight, mean=0, std=init_method_std * ( |
| 249 | (module.hidden_size * module.hidden_size_per_attention_head) ** -0.5)) |
| 250 | elif name == "key_value": |