| 28 | return use_batch_major_op_cache, x |
| 29 | |
| 30 | class FtDecoderWeights(object): |
| 31 | def __init__(self, layer_num, hidden_dim, onmtcheckpoint, max_step_for_pe=2048): |
| 32 | self.max_step_for_pe = max_step_for_pe |
| 33 | self.hidden_dim = hidden_dim |
| 34 | self.w = [] |
| 35 | prefix = 'decoder.transformer_layers.' |
| 36 | self.w.append(torch.stack( |
| 37 | [onmtcheckpoint['model'][prefix + str(i) + '.layer_norm_1.weight'] for i in range(layer_num)], |
| 38 | 0).contiguous()) |
| 39 | self.w.append(torch.stack( |
| 40 | [onmtcheckpoint['model'][prefix + str(i) + '.layer_norm_1.bias'] for i in range(layer_num)], |
| 41 | 0).contiguous()) |
| 42 | self.w.append(torch.stack( |
| 43 | [torch.stack([onmtcheckpoint['model'][prefix + str(i) + '.self_attn.linear_query.weight'].transpose(-1, -2), |
| 44 | onmtcheckpoint['model'][prefix + str(i) + '.self_attn.linear_keys.weight'].transpose(-1, -2), |
| 45 | onmtcheckpoint['model'][prefix + str(i) + '.self_attn.linear_values.weight'].transpose(-1, -2)], -2) |
| 46 | for i in range(layer_num)], 0).contiguous()) |
| 47 | self.w.append(torch.stack( |
| 48 | [torch.stack([onmtcheckpoint['model'][prefix + str(i) + '.self_attn.linear_query.bias'], |
| 49 | onmtcheckpoint['model'][prefix + str(i) + '.self_attn.linear_keys.bias'], |
| 50 | onmtcheckpoint['model'][prefix + str(i) + '.self_attn.linear_values.bias']], -2) for i in range(layer_num)], |
| 51 | 0).contiguous()) |
| 52 | self.w.append(torch.stack( |
| 53 | [onmtcheckpoint['model'][prefix + str(i) + '.self_attn.final_linear.weight'].transpose(-1, -2) for i in range(layer_num)], |
| 54 | 0).contiguous()) |
| 55 | self.w.append(torch.stack( |
| 56 | [onmtcheckpoint['model'][prefix + str(i) + '.self_attn.final_linear.bias'] for i in range(layer_num)], |
| 57 | 0).contiguous()) |
| 58 | self.w.append(torch.stack( |
| 59 | [onmtcheckpoint['model'][prefix + str(i) + '.layer_norm_2.weight'] for i in range(layer_num)], |
| 60 | 0).contiguous()) |
| 61 | self.w.append(torch.stack( |
| 62 | [onmtcheckpoint['model'][prefix + str(i) + '.layer_norm_2.bias'] for i in range(layer_num)], |
| 63 | 0).contiguous()) |
| 64 | self.w.append(torch.stack( |
| 65 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.linear_query.weight'].transpose(-1, -2) for i in range(layer_num)], |
| 66 | 0).contiguous()) |
| 67 | self.w.append(torch.stack( |
| 68 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.linear_keys.weight'].transpose(-1, -2) for i in range(layer_num)], |
| 69 | 0).contiguous()) |
| 70 | self.w.append(torch.stack( |
| 71 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.linear_values.weight'].transpose(-1, -2) for i in range(layer_num)], |
| 72 | 0).contiguous()) |
| 73 | self.w.append(torch.stack( |
| 74 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.linear_query.bias'] for i in range(layer_num)], |
| 75 | 0).contiguous()) |
| 76 | self.w.append(torch.stack( |
| 77 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.linear_keys.bias'] for i in range(layer_num)], |
| 78 | 0).contiguous()) |
| 79 | self.w.append(torch.stack( |
| 80 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.linear_values.bias'] for i in range(layer_num)], |
| 81 | 0).contiguous()) |
| 82 | self.w.append(torch.stack( |
| 83 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.final_linear.weight'].transpose(-1, -2) for i in range(layer_num)], |
| 84 | 0).contiguous()) |
| 85 | self.w.append(torch.stack( |
| 86 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.final_linear.bias'] for i in range(layer_num)], |
| 87 | 0).contiguous()) |