| 117 | self.w[i] = self.w[i].bfloat16() |
| 118 | |
| 119 | class FTDecoder(nn.Module): |
| 120 | def __init__(self, head_num, head_size, mem_hidden_dim, layer_num, weights, args): |
| 121 | super().__init__() |
| 122 | self.args = args |
| 123 | self.is_fp16 = True if self.args.data_type == 'fp16' else False |
| 124 | self.layer_num = layer_num |
| 125 | self.use_batch_major_op_cache, self.op_cache_dim_x = get_op_cache_config(head_size, self.is_fp16) |
| 126 | |
| 127 | torch.classes.load_library(args.decoder_ths_path) |
| 128 | try: |
| 129 | self.dec_layer = torch.classes.FasterTransformer.Decoder(*weights.w, head_num, head_size, head_num * head_size * 4, layer_num, mem_hidden_dim) |
| 130 | except: |
| 131 | # legacy ths for 20.03 image |
| 132 | self.dec_layer = torch.classes.FasterTransformerDecoder(*weights.w, head_num, head_size, head_num * head_size * 4, layer_num, mem_hidden_dim) |
| 133 | |
| 134 | def forward(self, inputs, memory, memory_seq_lens, self_cache, mem_cache, sequence_lengths, step): |
| 135 | dtype = torch.half if self.is_fp16 else torch.float32 |
| 136 | inputs_shape = inputs.shape |
| 137 | inputs = inputs.reshape([-1, inputs.shape[-1]]) |
| 138 | output, self_key_cache, self_val_cache, mem_key_cache, mem_val_cache = \ |
| 139 | self.dec_layer.forward(step, inputs, memory, memory_seq_lens, sequence_lengths, |
| 140 | self_cache[0], self_cache[1], mem_cache[0], mem_cache[1]) |
| 141 | output = output.reshape(inputs_shape) |
| 142 | |
| 143 | return output, [self_key_cache, self_val_cache], [mem_key_cache, mem_val_cache] |
| 144 | |
| 145 | |
| 146 | class ArgHelper(object): |