(self,
context_lengths,
batch_size,
beam_width,
scfg,
in_progress=False)
| 3112 | return final_output_ids |
| 3113 | |
| 3114 | def finalize_decoder(self, |
| 3115 | context_lengths, |
| 3116 | batch_size, |
| 3117 | beam_width, |
| 3118 | scfg, |
| 3119 | in_progress=False): |
| 3120 | final_output_ids = None |
| 3121 | if self.mapping.is_last_pp_rank(): |
| 3122 | # output shape of self.gather_tree: [batch_size, beam_width, output_len] |
| 3123 | beam_hyps_args = [ |
| 3124 | self.beam_hyps_output_ids_cba, self.beam_hyps_seq_len_cba, |
| 3125 | self.beam_hyps_cum_log_probs_cba, |
| 3126 | self.beam_hyps_normed_scores_cba, self.beam_hyps_log_probs_cba, |
| 3127 | self.beam_hyps_min_normed_scores, self.beam_hyps_num_beams, |
| 3128 | self.beam_hyps_is_done |
| 3129 | ] |
| 3130 | |
| 3131 | if scfg.use_beam_hyps and in_progress: |
| 3132 | # self.gather_tree modifies these args. |
| 3133 | # In streaming mode, this results in incorrect decoding in the following steps. |
| 3134 | beam_hyps_args = copy.deepcopy(beam_hyps_args) |
| 3135 | |
| 3136 | final_output_ids = self.gather_tree( |
| 3137 | self.sequence_length_buffer, self.output_ids, self.parent_ids, |
| 3138 | self.end_ids, context_lengths, self.cum_log_probs, |
| 3139 | self.log_probs, self.log_probs_tiled, *beam_hyps_args, |
| 3140 | self.finished, self.length_penalty, batch_size, beam_width, |
| 3141 | self.max_seq_length, scfg.use_beam_hyps) |
| 3142 | |
| 3143 | # Communicate ranks in Pipeline Parallelism |
| 3144 | if self.mapping.has_pp(): |
| 3145 | final_output_ids = self.pp_communicate_final_output_ids( |
| 3146 | final_output_ids, batch_size, beam_width) |
| 3147 | |
| 3148 | return final_output_ids |
| 3149 | |
| 3150 | def find_best_medusa_path(self, |
| 3151 | batch_size, |
no test coverage detected