| 421 | |
| 422 | |
| 423 | class TorchDecoding(nn.Module): |
| 424 | def __init__(self, layer_num, head_num, head_size, vocab_size, start_id, end_id, weights, |
| 425 | beam_search_diversity_rate=0.0, args=None): |
| 426 | super().__init__() |
| 427 | self.layer_num = layer_num |
| 428 | self.hidden_dim = head_num * head_size |
| 429 | self.start_id = start_id |
| 430 | self.end_id = end_id |
| 431 | self.vocab_size = vocab_size |
| 432 | self.diversity_rate = beam_search_diversity_rate |
| 433 | self.args = args |
| 434 | emb = Embeddings(self.hidden_dim, vocab_size, 1, position_encoding=True) |
| 435 | self.decoder = TransformerDecoder(layer_num, self.hidden_dim, head_num, head_size, 4*self.hidden_dim, |
| 436 | False, 'scaled-dot', 0, 0, emb, 0, False, False, -3, 0, args) |
| 437 | self.generator = nn.Linear(self.hidden_dim, vocab_size) |
| 438 | self.logsoftmax = nn.LogSoftmax(dim=-1) |
| 439 | if args.model_type == 'torch_decoding': |
| 440 | for i in range(layer_num): |
| 441 | prefix = 'decoder.transformer_layers.' + str(i) |
| 442 | self.decoder.transformer_layers[i].layer_norm_1.weight.data = weights.w['model'][prefix + '.layer_norm_1.weight'] |
| 443 | self.decoder.transformer_layers[i].layer_norm_1.bias.data = weights.w['model'][prefix + '.layer_norm_1.bias'] |
| 444 | self.decoder.transformer_layers[i].self_attn.linear_query.weight.data = weights.w['model'][prefix + '.self_attn.linear_query.weight'] |
| 445 | self.decoder.transformer_layers[i].self_attn.linear_keys.weight.data = weights.w['model'][prefix + '.self_attn.linear_keys.weight'] |
| 446 | self.decoder.transformer_layers[i].self_attn.linear_values.weight.data = weights.w['model'][prefix + '.self_attn.linear_values.weight'] |
| 447 | self.decoder.transformer_layers[i].self_attn.linear_query.bias.data = weights.w['model'][prefix + '.self_attn.linear_query.bias'] |
| 448 | self.decoder.transformer_layers[i].self_attn.linear_keys.bias.data = weights.w['model'][prefix + '.self_attn.linear_keys.bias'] |
| 449 | self.decoder.transformer_layers[i].self_attn.linear_values.bias.data = weights.w['model'][prefix + '.self_attn.linear_values.bias'] |
| 450 | self.decoder.transformer_layers[i].self_attn.final_linear.weight.data = weights.w['model'][prefix + '.self_attn.final_linear.weight'] |
| 451 | self.decoder.transformer_layers[i].self_attn.final_linear.bias.data = weights.w['model'][prefix + '.self_attn.final_linear.bias'] |
| 452 | self.decoder.transformer_layers[i].layer_norm_2.weight.data = weights.w['model'][prefix + '.layer_norm_2.weight'] |
| 453 | self.decoder.transformer_layers[i].layer_norm_2.bias.data = weights.w['model'][prefix + '.layer_norm_2.bias'] |
| 454 | self.decoder.transformer_layers[i].context_attn.linear_query.weight.data = weights.w['model'][prefix + '.context_attn.linear_query.weight'] |
| 455 | self.decoder.transformer_layers[i].context_attn.linear_keys.weight.data = weights.w['model'][prefix + '.context_attn.linear_keys.weight'] |
| 456 | self.decoder.transformer_layers[i].context_attn.linear_values.weight.data = weights.w['model'][prefix + '.context_attn.linear_values.weight'] |
| 457 | self.decoder.transformer_layers[i].context_attn.linear_query.bias.data = weights.w['model'][prefix + '.context_attn.linear_query.bias'] |
| 458 | self.decoder.transformer_layers[i].context_attn.linear_keys.bias.data = weights.w['model'][prefix + '.context_attn.linear_keys.bias'] |
| 459 | self.decoder.transformer_layers[i].context_attn.linear_values.bias.data = weights.w['model'][prefix + '.context_attn.linear_values.bias'] |
| 460 | self.decoder.transformer_layers[i].context_attn.final_linear.weight.data = weights.w['model'][prefix + '.context_attn.final_linear.weight'] |
| 461 | self.decoder.transformer_layers[i].context_attn.final_linear.bias.data = weights.w['model'][prefix + '.context_attn.final_linear.bias'] |
| 462 | self.decoder.transformer_layers[i].feed_forward.layer_norm.weight.data = weights.w['model'][prefix + '.feed_forward.layer_norm.weight'] |
| 463 | self.decoder.transformer_layers[i].feed_forward.layer_norm.bias.data = weights.w['model'][prefix + '.feed_forward.layer_norm.bias'] |
| 464 | self.decoder.transformer_layers[i].feed_forward.w_1.weight.data = weights.w['model'][prefix + '.feed_forward.w_1.weight'] |
| 465 | self.decoder.transformer_layers[i].feed_forward.w_1.bias.data = weights.w['model'][prefix + '.feed_forward.w_1.bias'] |
| 466 | self.decoder.transformer_layers[i].feed_forward.w_2.weight.data = weights.w['model'][prefix + '.feed_forward.w_2.weight'] |
| 467 | self.decoder.transformer_layers[i].feed_forward.w_2.bias.data = weights.w['model'][prefix + '.feed_forward.w_2.bias'] |
| 468 | elif args.model_type == 'torch_decoding_with_decoder_ext': |
| 469 | w = [] |
| 470 | ft_decoder_weights = FtDecoderWeights(layer_num, self.hidden_dim, weights.w) |
| 471 | ft_decoder_weights.to_cuda() |
| 472 | if args.data_type == 'fp16': |
| 473 | ft_decoder_weights.to_half() |
| 474 | elif args.data_type == 'bf16': |
| 475 | ft_decoder_weights.to_bfloat16() |
| 476 | self.decoder.transformer_layers = nn.ModuleList( |
| 477 | [FTDecoder(head_num, head_size, head_num * head_size, layer_num, ft_decoder_weights, args)]) |
| 478 | else: |
| 479 | raise ValueError('wrong model_type') |
| 480 | self.decoder.layer_norm.weight.data = weights.w['model']['decoder.layer_norm.weight'] |
no outgoing calls
no test coverage detected