Add the transformer layer
(prefix, config, init_dict, network, input_tensor, residual, mask_idx, cu_seqlens, max_seqlen)
| 185 | return layer |
| 186 | |
| 187 | def transformer_layer_opt(prefix, config, init_dict, network, input_tensor, residual, mask_idx, cu_seqlens, max_seqlen): |
| 188 | """ |
| 189 | Add the transformer layer |
| 190 | """ |
| 191 | hidden_size = config.hidden_size |
| 192 | |
| 193 | if config.use_qat: |
| 194 | dr_input = init_dict[prefix + 'attention_self_query_input_amax'] |
| 195 | assert(dr_input ==init_dict[prefix + 'attention_self_key_input_amax'] ) |
| 196 | assert(dr_input ==init_dict[prefix + 'attention_self_value_input_amax'] ) |
| 197 | input_tensor.set_dynamic_range(-dr_input, dr_input) |
| 198 | |
| 199 | context_transposed = attention_layer_opt(prefix + "attention_", config, init_dict, network, input_tensor, mask_idx, cu_seqlens, max_seqlen) |
| 200 | attention_heads = context_transposed.get_output(0) |
| 201 | |
| 202 | # FC0 |
| 203 | B_aout = init_dict[prefix + B_AOUT] |
| 204 | W_aout = init_dict[prefix + W_AOUT] |
| 205 | if config.use_int8: |
| 206 | attention_out_fc = network.add_convolution_nd(attention_heads, hidden_size, (1, 1), W_aout, B_aout) |
| 207 | else: |
| 208 | attention_out_fc = network.add_fully_connected(attention_heads, hidden_size, W_aout, B_aout) |
| 209 | if config.use_int8 and config.use_qat: |
| 210 | dr_fc_aout = init_dict[prefix + 'attention_output_add_local_input_quantizer_amax'] |
| 211 | set_output_range(attention_out_fc, dr_fc_aout) |
| 212 | |
| 213 | if config.use_megatron: |
| 214 | dr_skln1_res_in = init_dict[prefix + "attention_output_add_residual_input_quantizer_amax"] |
| 215 | residual.set_dynamic_range(-dr_skln1_res_in, dr_skln1_res_in) |
| 216 | skip = residual |
| 217 | else: |
| 218 | skip = input_tensor |
| 219 | skiplayer = skipln(prefix + "attention_output_layernorm_", config, init_dict, network, attention_out_fc.get_output(0), skip) |
| 220 | attention_ln = skiplayer.get_output(0) |
| 221 | if config.use_qat: |
| 222 | dr_skln1 = init_dict[prefix + 'intermediate_dense_input_amax'] |
| 223 | set_output_range(skiplayer, dr_skln1) |
| 224 | |
| 225 | # FC1 + GELU |
| 226 | B_mid = init_dict[prefix + B_MID] |
| 227 | W_mid = init_dict[prefix + W_MID] |
| 228 | if config.use_int8: |
| 229 | mid_dense = network.add_convolution_nd(attention_ln, config.intermediate_size, (1, 1), W_mid, B_mid) |
| 230 | else: |
| 231 | mid_dense = network.add_fully_connected(attention_ln, config.intermediate_size, W_mid, B_mid) |
| 232 | |
| 233 | gelu_layer = add_gelu(network, mid_dense.get_output(0)) |
| 234 | |
| 235 | intermediate_act = gelu_layer.get_output(0) |
| 236 | set_tensor_name(intermediate_act, prefix, "gelu") |
| 237 | if config.use_int8: |
| 238 | if config.use_qat: |
| 239 | dr_gelu = init_dict[prefix + 'output_dense_input_amax'] |
| 240 | set_output_range(gelu_layer, dr_gelu) |
| 241 | else: |
| 242 | # use gelu10 according to whitepaper http://arxiv.org/abs/2004.09602 |
| 243 | set_output_range(gelu_layer, 10) |
| 244 |
no test coverage detected