def post_language_model_processing( lm_output, labels, logit_weights, get_key_value, parallel_output, forward_method_parallel_output, fp16_lm_cross_entropy, return_logits=False, ): if get_key_value: lm_output, presents = lm_output # Output. if forward_method_parallel_output is not None: parallel_output = forward_method_parallel_output output = parallel_lm_logits(lm_output, logit_weights, parallel_output) if get_key_value: output = [output, presents] if labels is None: return output else: if fp16_lm_cross_entropy: assert output.dtype == torch.half loss = tensor_parallel.vocab_parallel_cross_entropy(output, labels) else: loss = tensor_parallel.vocab_parallel_cross_entropy( output.float(), labels) if return_logits: return loss, output else: return loss
def forward( self, encoder_input_ids, decoder_input_ids, encoder_attn_mask, decoder_attn_mask, encoder_decoder_attn_mask, tokentype_ids=None, lm_labels=None, enc_hidden_states=None, output_enc_hidden_only=False, ): # Converting the attention masks to proper parameter settings encoder_attn_mask, decoder_attn_mask, encoder_decoder_attn_mask = t5_extended_attention_mask( [encoder_attn_mask, decoder_attn_mask, encoder_decoder_attn_mask]) encoder_position_ids = t5_position_ids(encoder_input_ids) # Handle case when decoder_input_ids is None to get just the encoder hidden states. decoder_position_ids = t5_position_ids( decoder_input_ids) if decoder_input_ids is not None else None lm_output = self.language_model( enc_input_ids=encoder_input_ids, enc_position_ids=encoder_position_ids, enc_attn_mask=encoder_attn_mask, dec_input_ids=decoder_input_ids, dec_position_ids=decoder_position_ids, dec_attn_mask=decoder_attn_mask, enc_dec_attn_mask=encoder_decoder_attn_mask, tokentype_ids=tokentype_ids, enc_hidden_states=enc_hidden_states, output_enc_hidden_only=output_enc_hidden_only, ) if output_enc_hidden_only: return lm_output decoder_output, encoder_output = lm_output # Output. lm_logits = self.lm_head( decoder_output, self.language_model.embedding.word_embeddings.weight) if lm_labels is None: return lm_logits, encoder_output else: if self.fp16_lm_cross_entropy: assert lm_logits.dtype == torch.half lm_loss = tensor_parallel.vocab_parallel_cross_entropy( lm_logits, lm_labels) else: lm_loss = tensor_parallel.vocab_parallel_cross_entropy( lm_logits.float(), lm_labels) return lm_loss, encoder_output
def forward( self, input_ids, input_attn_mask, retrieved_emb, retrieved_attn_mask, token_type_ids=None, labels=None, input_emb=None, ): """ Return value is per token / per dimension (i.e., non collapsed loss value) """ if input_emb is None: if self.pre_process and self.add_encoder: # encoder embeddings input_position_ids = build_position_ids(input_ids) input_emb = self.encoder_embedding( input_ids, input_position_ids, token_type_ids=token_type_ids) else: input_emb = None if self.add_decoder: hidden = self.pre_decoder(input_emb, input_attn_mask) if self.add_encoder: retrieved_emb = self.encoder(retrieved_emb, retrieved_attn_mask, context_attn_mask=input_attn_mask, encoder_output=hidden) if self.add_decoder: dec_output = self.post_decoder( hidden, input_attn_mask, retrieved_attn_mask=retrieved_attn_mask, retrieved_emb=retrieved_emb) token_logits = self.tokens_head(dec_output, self.word_embeddings_weight()) if labels is not None: # tensor_parallel.vocab_parallel_cross_entropy performs log_softmax and return log p(x_i|z) per token i if self.fp16_cross_entropy: assert token_logits.dtype == torch.half tokens_loss = tensor_parallel.vocab_parallel_cross_entropy( token_logits, labels) else: tokens_loss = tensor_parallel.vocab_parallel_cross_entropy( token_logits.float(), labels) return tokens_loss else: return token_logits
def post_language_model_processing( lm_output, pooled_output, lm_head, binary_head, lm_labels, logit_weights, fp16_lm_cross_entropy ): # Output. lm_logits = lm_head(lm_output, logit_weights) binary_logits = None if binary_head is not None: binary_logits = binary_head(pooled_output) if lm_labels is None: return lm_logits, binary_logits else: if fp16_lm_cross_entropy: assert lm_logits.dtype == torch.half lm_loss = tensor_parallel.vocab_parallel_cross_entropy(lm_logits, lm_labels) else: lm_loss = tensor_parallel.vocab_parallel_cross_entropy(lm_logits.float(), lm_labels) return lm_loss, binary_logits
def forward( self, input_ids, input_attn_mask, retrieved_ids, retrieved_attn_mask, token_type_ids=None, labels=None, input_emb=None, set_inference_key_value_memory=False, inference_max_sequence_len=None, neighbors=None, ): """ Return value is per token / per dimension (i.e., non collapsed loss value) """ eod_positions = None retrieved_emb = None if input_ids is not None and self.eod_id is not None: eod_positions = torch.where(input_ids == self.eod_id) if input_emb is None: if self.pre_process and self.add_encoder: # encoder embeddings if self.add_abs_position_embedding: input_position_ids = build_position_ids(input_ids) else: input_position_ids = None input_emb = self.encoder_embedding(input_ids, input_position_ids, token_type_ids=token_type_ids) else: input_emb = None if retrieved_ids is not None: if self.add_abs_position_embedding: seq_length = retrieved_ids.size(-1) retrieved_position_ids = torch.arange(seq_length, dtype=torch.long, device=retrieved_ids.device) retrieved_position_ids = retrieved_position_ids.unsqueeze(0).expand_as(retrieved_ids).clone() else: retrieved_position_ids = None retrieved_emb = self.encoder_embedding(retrieved_ids, retrieved_position_ids) if self.add_decoder: hidden = self.pre_decoder( input_emb, input_attn_mask, eod_positions=eod_positions, set_inference_key_value_memory=set_inference_key_value_memory, inference_max_sequence_len=inference_max_sequence_len, ) # hidden is a tuple, (layernorm_input, layernorm_output) self.post_decoder.set_input_tensor(hidden) encoder_input = hidden[1].transpose(0, 1).contiguous() if self.add_encoder: if retrieved_emb is not None and neighbors is None: neighbors = retrieved_emb.shape[2] retrieved_emb = self.encoder( retrieved_emb, retrieved_attn_mask, context_attn_mask=input_attn_mask, encoder_output=encoder_input, set_inference_key_value_memory=set_inference_key_value_memory, inference_max_sequence_len=inference_max_sequence_len, neighbors=neighbors, ) if self.add_decoder: dec_output = self.post_decoder( hidden, input_attn_mask, retrieved_attn_mask=retrieved_attn_mask, retrieved_emb=retrieved_emb, eod_positions=eod_positions, set_inference_key_value_memory=set_inference_key_value_memory, inference_max_sequence_len=inference_max_sequence_len, ) token_logits = self.tokens_head(dec_output, self.word_embeddings_weight()) if labels is not None: # tensor_parallel.vocab_parallel_cross_entropy performs log_softmax and return log p(x_i|z) per token i if self.fp16_cross_entropy: assert token_logits.dtype == torch.half tokens_loss = tensor_parallel.vocab_parallel_cross_entropy(token_logits, labels) else: tokens_loss = tensor_parallel.vocab_parallel_cross_entropy(token_logits.float(), labels) return tokens_loss else: return token_logits
def forward( self, enc_input_ids, enc_attn_mask, dec_input_ids, dec_attn_mask, token_type_ids=None, labels=None, enc_hidden_states=None, enc_output_mask=None, output_enc_hidden_only=False, enc_input=None, ): """ Return value is per token / per dimension (i.e., non collapsed loss value) """ if enc_input is None: if self.pre_process and self.add_encoder: # encoder embeddings enc_position_ids = build_position_ids(enc_input_ids) enc_input = self.encoder_embedding( enc_input_ids, enc_position_ids, token_type_ids=token_type_ids) else: enc_input = None if output_enc_hidden_only: enc_output = self.enc_dec_model.encode( enc_input=enc_input, enc_attn_mask=enc_attn_mask, enc_layer_past=None, enc_get_key_value=False, ) return enc_output else: if self.pre_process and self.add_decoder: dec_position_ids = build_position_ids(dec_input_ids) dec_input = self.decoder_embedding( dec_input_ids, dec_position_ids, token_type_ids=token_type_ids) else: # Note: This is when the decoder itself is split across PP ranks. dec_input = None output = self.enc_dec_model( enc_input=enc_input, enc_attn_mask=enc_attn_mask, dec_input=dec_input, dec_attn_mask=dec_attn_mask, enc_layer_past=None, enc_get_key_value=False, enc_output=None, dec_layer_past=None, dec_get_key_value=False, ) if self.post_process and self.add_decoder: dec_output, enc_output = output # project decoder output to vocabulary-size dimensions token_logits = self.tokens_head(dec_output, self.word_embeddings_weight()) if labels is not None: # tensor_parallel.vocab_parallel_cross_entropy performs log_softmax and return log p(x_i|z) per token i if self.fp16_cross_entropy: assert token_logits.dtype == torch.half tokens_loss = tensor_parallel.vocab_parallel_cross_entropy( token_logits, labels) else: tokens_loss = tensor_parallel.vocab_parallel_cross_entropy( token_logits.float(), labels) return tokens_loss else: return token_logits elif self.add_decoder and not self.add_encoder: decoder_output, _ = output return decoder_output else: encoder_output = output return encoder_output
def forward( self, enc_input_ids, enc_attn_mask, dec_input_ids, dec_attn_mask, tokentype_ids=None, labels=None, enc_hidden_states=None, enc_output_mask=None, output_enc_hidden_only=False, enc_input=None, ): """ Return value is per token / per dimension (i.e., non collapsed loss value) """ ret_dict = {} # encoder embeddings if enc_input is None: enc_position_ids = build_position_ids(enc_input_ids) enc_input = self.encoder_embedding(enc_input_ids, enc_position_ids, tokentype_ids=tokentype_ids) if output_enc_hidden_only: enc_output, enc_output_mask = self.enc_dec_model.encode( enc_input=enc_input, enc_attn_mask=enc_attn_mask, enc_layer_past=None, enc_get_key_value=False, ) ret_dict["enc_output"] = enc_output ret_dict["enc_output_mask"] = enc_output_mask else: dec_position_ids = build_position_ids(dec_input_ids) dec_input = self.decoder_embedding(dec_input_ids, dec_position_ids, tokentype_ids=tokentype_ids) ret_dict.update( self.enc_dec_model( enc_input=enc_input, enc_attn_mask=enc_attn_mask, dec_input=dec_input, dec_attn_mask=dec_attn_mask, enc_layer_past=None, enc_get_key_value=False, enc_output=enc_hidden_states, enc_output_mask=enc_output_mask, dec_layer_past=None, dec_get_key_value=False, )) # project decoder output to vocabulary-size dimensions token_logits = self.tokens_head( ret_dict["dec_output"], self.decoder_embedding.word_embeddings.weight) # token_logits [batch, length, vocab_size] ret_dict["token_logits"] = token_logits if labels is not None: # tensor_parallel.vocab_parallel_cross_entropy performs log_softmax and return log p(x_i|z) per token i if self.fp16_cross_entropy: assert token_logits.dtype == torch.half tokens_loss = tensor_parallel.vocab_parallel_cross_entropy( token_logits, labels) else: tokens_loss = tensor_parallel.vocab_parallel_cross_entropy( token_logits.float(), labels) # tokens_loss [batch, length] ret_dict["tokens_loss"] = tokens_loss return ret_dict