(batch_sizes, workspace_size, sequence_length, config, weights_dict, squad_json, vocab_file, calibrationCacheFile, calib_num, verbose)
| 394 | return emb_layer, cu_seqlens, max_seqlen |
| 395 | |
| 396 | def build_engine(batch_sizes, workspace_size, sequence_length, config, weights_dict, squad_json, vocab_file, calibrationCacheFile, calib_num, verbose): |
| 397 | explicit_batch_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) |
| 398 | |
| 399 | with trt.Builder(TRT_LOGGER) as builder, builder.create_network(explicit_batch_flag) as network, builder.create_builder_config() as builder_config: |
| 400 | builder_config.max_workspace_size = workspace_size * (1024 * 1024) |
| 401 | builder_config.avg_timing_iterations = 8 |
| 402 | if config.use_fp16: |
| 403 | builder_config.set_flag(trt.BuilderFlag.FP16) |
| 404 | if config.use_int8: |
| 405 | builder_config.set_flag(trt.BuilderFlag.INT8) |
| 406 | if not config.use_qat: |
| 407 | raise RuntimeError("Post training calibration is not supported in variable-length BERT.") |
| 408 | |
| 409 | if verbose: |
| 410 | builder_config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED |
| 411 | |
| 412 | # speed up the engine build for trt major version >= 8 |
| 413 | # 1. disable cudnn tactic |
| 414 | # 2. load global timing cache |
| 415 | if trt_version[0] >= 8: |
| 416 | tactic_source = builder_config.get_tactic_sources() & ~(1 << int(trt.TacticSource.CUDNN)) |
| 417 | builder_config.set_tactic_sources(tactic_source) |
| 418 | if config.timing_cache != None: |
| 419 | if os.path.exists(config.timing_cache): |
| 420 | with open(config.timing_cache, "rb") as f: |
| 421 | cache = builder_config.create_timing_cache(f.read()) |
| 422 | builder_config.set_timing_cache(cache, ignore_mismatch = False) |
| 423 | else: |
| 424 | cache = builder_config.create_timing_cache(b"") |
| 425 | builder_config.set_timing_cache(cache, ignore_mismatch = False) |
| 426 | |
| 427 | if config.use_sparsity: |
| 428 | TRT_LOGGER.log(TRT_LOGGER.INFO, "Setting sparsity flag on builder_config.") |
| 429 | builder_config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS) |
| 430 | |
| 431 | # Create the network |
| 432 | emb_layer, cu_seqlens, max_seqlen = emb_layernorm(builder, network, config, weights_dict, builder_config, sequence_length, batch_sizes) |
| 433 | embeddings = emb_layer.get_output(0) |
| 434 | if config.use_int8 and config.interleaved: |
| 435 | shuffle = network.add_shuffle(embeddings) |
| 436 | shuffle.second_transpose = (2, 1, 0, 3) |
| 437 | embeddings = shuffle.get_output(0) |
| 438 | mask_idx = None |
| 439 | else: |
| 440 | mask_idx = emb_layer.get_output(1) |
| 441 | |
| 442 | if config.use_megatron: # megatron currently only supports int8 and interleaved |
| 443 | shuffler = network.add_shuffle(emb_layer.get_output(1)) |
| 444 | shuffler.second_transpose = (2, 1, 0, 3) |
| 445 | residual = shuffler.get_output(0) |
| 446 | |
| 447 | dr_emb = weights_dict['l0_attention_self_query_input_amax'] |
| 448 | embeddings.set_dynamic_range(-dr_emb, dr_emb) |
| 449 | dr_skln1_res_in = weights_dict['l0_attention_output_add_residual_input_quantizer_amax'] |
| 450 | residual.set_dynamic_range(-dr_skln1_res_in, dr_skln1_res_in) |
| 451 | else: |
| 452 | residual = None |
| 453 |
no test coverage detected