wehub-resource-sync
197 lines
7.4 KiB
Python
197 lines
7.4 KiB
Python
import torch
|
|
from torch import nn
|
|
from transformers.models.llama.modeling_llama import (
|
|
LlamaForCausalLM,
|
|
LlamaModel,
|
|
)
|
|
from transformers.models.gemma.modeling_gemma import (
|
|
GemmaForCausalLM,
|
|
GemmaModel,
|
|
)
|
|
from transformers.models.mistral.modeling_mistral import (
|
|
MistralForCausalLM,
|
|
MistralPreTrainedModel,
|
|
)
|
|
from typing import Union
|
|
|
|
|
|
def sft_loss_on_logits(logits: torch.FloatTensor, labels: torch.LongTensor, pad_token_id: int, macro_average: bool = False, row_weights: torch.Tensor = None):
|
|
batch_size = labels.size(0)
|
|
labels = labels[:, 1:].contiguous()
|
|
logits = logits[:, :-1].contiguous()
|
|
|
|
logits = logits.view(-1, logits.size(-1))
|
|
labels = labels.view(-1)
|
|
|
|
pad_mask = labels.eq(pad_token_id)
|
|
if pad_mask.sum() == labels.numel(): # To tackle all problems that the response are empty, or the pad_token equals eos_token so no response.
|
|
return 0.
|
|
|
|
labels[pad_mask] = -100
|
|
if macro_average:
|
|
loss_fct = nn.CrossEntropyLoss(reduction='none')
|
|
loss = loss_fct(logits, labels)
|
|
|
|
row_num_element = (~pad_mask).reshape(batch_size, -1).sum(-1).float()
|
|
row_mask = row_num_element > 0
|
|
loss = loss.view(batch_size, -1)
|
|
loss = loss.sum(-1) / row_num_element
|
|
if row_weights is not None:
|
|
loss = loss * row_weights
|
|
loss = loss[row_mask].mean()
|
|
else:
|
|
loss_fct = nn.CrossEntropyLoss()
|
|
loss = loss_fct(logits, labels)
|
|
return loss
|
|
|
|
|
|
def llama_dpo_batch_forward(model: Union[LlamaForCausalLM, GemmaForCausalLM, MistralForCausalLM],
|
|
input_ids: torch.LongTensor, attention_mask: torch.Tensor, labels: torch.LongTensor, pad_token_id: int = None,
|
|
average_log_prob: bool = False):
|
|
outputs = model.model(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
)
|
|
|
|
hidden_states = outputs[0]
|
|
logits = model.lm_head(hidden_states)
|
|
logits = logits.float()
|
|
|
|
labels = labels[:, 1:].clone()
|
|
|
|
if pad_token_id is None:
|
|
pad_token_id = model.config.pad_token_id
|
|
|
|
loss_mask = labels.ne(pad_token_id)
|
|
labels[~loss_mask] = 0
|
|
|
|
per_token_logprobs = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)
|
|
if average_log_prob:
|
|
log_ps = (per_token_logprobs * loss_mask).sum(-1) / loss_mask.sum(-1)
|
|
else:
|
|
log_ps = (per_token_logprobs * loss_mask).sum(-1)
|
|
|
|
return logits, log_ps, loss_mask
|
|
|
|
|
|
def llama_batch_forward(model: Union[LlamaForCausalLM, GemmaForCausalLM, MistralForCausalLM], input_ids: torch.LongTensor, attention_mask: torch.Tensor):
|
|
outputs = model.model(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
)
|
|
|
|
hidden_states = outputs[0]
|
|
logits = model.lm_head(hidden_states)
|
|
|
|
return logits
|
|
|
|
|
|
def tdpo_get_batch_logps(logits: torch.FloatTensor, reference_logits: torch.FloatTensor, labels: torch.LongTensor, pad_token_id: int,
|
|
average_log_prob: bool = False):
|
|
"""Compute the kl divergence/log probabilities of the given labels under the given logits.
|
|
|
|
Args:
|
|
logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size)
|
|
reference_logits: Logits of the reference model (unnormalized). Shape: (batch_size, sequence_length, vocab_size)
|
|
labels: Labels for which to compute the log probabilities. Label tokens with a value of -100 are ignored. Shape: (batch_size, sequence_length)
|
|
pad_token_id: The id of the padding token.
|
|
average_log_prob: If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the log probabilities of the
|
|
(non-masked) tokens.
|
|
|
|
Returns:
|
|
Several tensors of shape (batch_size,) containing the average/sum kl divergence/log probabilities of the given labels under the given logits.
|
|
"""
|
|
assert logits.shape[:-1] == labels.shape
|
|
assert reference_logits.shape[:-1] == labels.shape
|
|
|
|
labels = labels[:, 1:].clone()
|
|
logits = logits[:, :-1, :]
|
|
reference_logits = reference_logits[:, :-1, :]
|
|
|
|
loss_mask = labels.ne(pad_token_id)
|
|
|
|
# dummy token; we'll ignore the losses on these tokens later
|
|
labels[~loss_mask] = 0
|
|
|
|
vocab_logps = logits.log_softmax(-1)
|
|
|
|
reference_vocab_ps = reference_logits.softmax(-1)
|
|
reference_vocab_logps = reference_vocab_ps.log()
|
|
|
|
per_position_kl = (reference_vocab_ps * (reference_vocab_logps - vocab_logps)).sum(-1)
|
|
per_token_logps = torch.gather(vocab_logps, dim=2, index=labels.unsqueeze(2)).squeeze(2)
|
|
per_reference_token_logps = torch.gather(reference_vocab_logps, dim=2, index=labels.unsqueeze(2)).squeeze(2)
|
|
|
|
logps_margin = per_token_logps - per_reference_token_logps
|
|
|
|
if average_log_prob:
|
|
return (logps_margin * loss_mask).sum(-1) / loss_mask.sum(-1), \
|
|
(per_position_kl * loss_mask).sum(-1) / loss_mask.sum(-1), \
|
|
(per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)
|
|
else:
|
|
return (logps_margin * loss_mask).sum(-1), \
|
|
(per_position_kl * loss_mask).sum(-1), \
|
|
(per_token_logps * loss_mask).sum(-1)
|
|
|
|
|
|
def llama_last_token_cls_batch_forward(model: Union[LlamaModel, GemmaForCausalLM, MistralPreTrainedModel], linear: nn.Linear,
|
|
input_ids: torch.LongTensor, attention_mask: torch.Tensor,
|
|
pad_token_id: int, return_full_logits: bool = False):
|
|
transformer_outputs = model(
|
|
input_ids,
|
|
attention_mask=attention_mask,
|
|
)
|
|
hidden_states = transformer_outputs[0]
|
|
|
|
batch_size = input_ids.shape[0]
|
|
sequence_lengths = (torch.eq(input_ids, pad_token_id).long().argmax(-1) - 1).to(device=hidden_states.device)
|
|
|
|
if return_full_logits:
|
|
rewards = linear(hidden_states)
|
|
return rewards, sequence_lengths
|
|
|
|
last_token_states = hidden_states[torch.arange(batch_size, device=hidden_states.device), sequence_lengths]
|
|
rewards = linear(last_token_states)
|
|
return rewards, sequence_lengths
|
|
|
|
|
|
def llama_token_batch_forward(model: Union[LlamaModel, GemmaModel], linear: nn.Linear,
|
|
input_ids: torch.LongTensor, attention_mask: torch.Tensor, pad_token_id: int = None, average: bool = False):
|
|
outputs = model(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
)
|
|
|
|
hidden_states = outputs[0]
|
|
logits = linear(hidden_states).squeeze(-1)
|
|
logits = logits.float()
|
|
|
|
if pad_token_id is None:
|
|
pad_token_id = model.config.pad_token_id
|
|
|
|
loss_mask = input_ids.ne(pad_token_id)
|
|
|
|
if average:
|
|
rewards = (logits * loss_mask).sum(-1) / loss_mask.sum(-1)
|
|
else:
|
|
rewards = (logits * loss_mask).sum(-1)
|
|
|
|
return rewards
|
|
|
|
|
|
def llama_last_token_forward_value(model: Union[LlamaModel, GemmaForCausalLM], linear: nn.Linear,
|
|
input_ids: torch.LongTensor, attention_mask: torch.Tensor,
|
|
pad_token_id: int):
|
|
transformer_outputs = model(
|
|
input_ids,
|
|
attention_mask=attention_mask,
|
|
)
|
|
hidden_states = transformer_outputs[0]
|
|
|
|
batch_size = input_ids.shape[0]
|
|
sequence_lengths = (torch.eq(input_ids, pad_token_id).long().argmax(-1) - 1).to(device=hidden_states.device)
|
|
|
|
values = linear(hidden_states)
|
|
rewards = values[torch.arange(batch_size, device=hidden_states.device), sequence_lengths]
|
|
return values, rewards, sequence_lengths
|