| 424 | |
| 425 | |
| 426 | class FTT5Decoding(nn.Module): |
| 427 | def __init__(self, decoding_weight_list, lib_path, head_num, head_size, inter_size, |
| 428 | mem_d_model, d_model, num_layer, start_id, end_id, vocab_size, q_scaling=1.0, num_bucket=32, |
| 429 | num_expert=0, moe_layer_index=[], |
| 430 | max_distance=128, tensor_para_size=1, pipeline_para_size=1, t5_with_bias=False, |
| 431 | position_embedding_type=0, moe_k=0, |
| 432 | activation_type="relu", tie_word_embeddings=True, adapter_inter_size=0, adapter_norm_position="pre"): |
| 433 | super().__init__() |
| 434 | |
| 435 | self.use_mpi = dist.is_mpi_available() |
| 436 | |
| 437 | if self.use_mpi: |
| 438 | try: |
| 439 | dist.init_process_group(backend='mpi') |
| 440 | except: |
| 441 | print("[INFO] WARNING: Exception occurred in dist.init_process_group(backend = 'mpi'). Maybe the process group has been initialized somewhere else.") |
| 442 | else: |
| 443 | print("[INFO] MPI is not available in this PyTorch build.") |
| 444 | assert tensor_para_size == 1, "[FATAL] MPI is required for tensor_para_size > 1." |
| 445 | assert pipeline_para_size == 1, "[FATAL] MPI is required for pipeline_para_size > 1." |
| 446 | |
| 447 | torch.classes.load_library(lib_path) |
| 448 | try: |
| 449 | self.decoding = torch.classes.FasterTransformer.T5Decoding(head_num, head_size, inter_size, mem_d_model, |
| 450 | d_model, num_layer, |
| 451 | vocab_size, num_bucket, num_expert, max_distance, |
| 452 | q_scaling, start_id, end_id, |
| 453 | tensor_para_size, pipeline_para_size, |
| 454 | t5_with_bias, |
| 455 | position_embedding_type, moe_k, activation_type, |
| 456 | tie_word_embeddings, adapter_inter_size, |
| 457 | adapter_norm_position, |
| 458 | moe_layer_index, *decoding_weight_list) |
| 459 | except: |
| 460 | self.decoding = torch.classes.FasterTransformerT5Decoding(head_num, head_size, inter_size, mem_d_model, |
| 461 | d_model, num_layer, |
| 462 | vocab_size, num_bucket, num_expert, max_distance, |
| 463 | q_scaling, start_id, end_id, |
| 464 | tensor_para_size, pipeline_para_size, |
| 465 | t5_with_bias, |
| 466 | position_embedding_type, moe_k, activation_type, |
| 467 | tie_word_embeddings, adapter_inter_size, |
| 468 | adapter_norm_position, |
| 469 | moe_layer_index, *decoding_weight_list) |
| 470 | |
| 471 | def forward(self, beam_width, max_seq_len, top_k, top_p, |
| 472 | beam_search_diversity_rate, temperature, |
| 473 | len_penalty, repetition_penalty, presence_penalty, min_length, random_seed, |
| 474 | mem_hidden_states, mem_seq_len, |
| 475 | is_return_output_log_probs, is_return_cum_log_probs, is_return_cross_attentions=False, |
| 476 | bad_words_list=None, stop_words_list=None): |
| 477 | # TODO (bhsueh) Not found an method to put a None Type into op forward function |
| 478 | # So, the top_k and top_p must be some values now. |
| 479 | results = self.decoding.forward(beam_width, max_seq_len, |
| 480 | top_k, top_p, beam_search_diversity_rate, |
| 481 | temperature, len_penalty, repetition_penalty, presence_penalty, min_length, |
| 482 | random_seed, mem_hidden_states, mem_seq_len, |
| 483 | is_return_output_log_probs, is_return_cum_log_probs, is_return_cross_attentions, |