chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,196 @@
|
||||
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
|
||||
@@ -0,0 +1,46 @@
|
||||
from transformers import PreTrainedModel
|
||||
import deepspeed
|
||||
from fairscale.nn.model_parallel import initialize as mpu
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
from general_util import training_utils
|
||||
|
||||
|
||||
def init_ds_training_engine(model: PreTrainedModel, ds_cfg: DictConfig, global_cfg: DictConfig, ):
|
||||
ds_config = ds_cfg
|
||||
if "total_num_steps" in ds_config.scheduler.params:
|
||||
ds_config.scheduler.params.total_num_steps = global_cfg.max_steps
|
||||
ds_config.scheduler.params.warmup_num_steps = global_cfg.warmup_steps
|
||||
ds_config = OmegaConf.to_container(ds_config, resolve=True)
|
||||
ds_config["train_mirco_batch_size_per_gpu"] = global_cfg.per_gpu_train_batch_size
|
||||
|
||||
optim_params = training_utils.get_optimizer_grouped_parameters(model, global_cfg.actor_weight_decay)
|
||||
|
||||
engine, optimizer, _, scheduler = deepspeed.initialize(
|
||||
model=model,
|
||||
model_parameters=optim_params,
|
||||
config_params=ds_config,
|
||||
mpu=mpu if mpu.model_parallel_is_initialized() else None,
|
||||
)
|
||||
|
||||
return engine, optimizer, scheduler
|
||||
|
||||
|
||||
def init_ds_eval_engine(model: PreTrainedModel, ds_cfg: DictConfig):
|
||||
ds_config = ds_cfg
|
||||
if ds_config.zero_optimization.stage != 3:
|
||||
ds_config.zero_optimization.stage = 0
|
||||
|
||||
ds_config = OmegaConf.to_container(ds_config, resolve=True)
|
||||
# ds_config["train_mirco_batch_size_per_gpu"] = global_cfg.per_gpu_train_batch_size
|
||||
if "optimizer" in ds_config:
|
||||
ds_config.pop("optimizer")
|
||||
if "scheduler" in ds_config:
|
||||
ds_config.pop("scheduler")
|
||||
|
||||
engine, *_ = deepspeed.initialize(
|
||||
model=model,
|
||||
config_params=ds_config,
|
||||
mpu=mpu if mpu.model_parallel_is_initialized() else None,
|
||||
)
|
||||
|
||||
return engine
|
||||
@@ -0,0 +1,45 @@
|
||||
import os
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
from fairscale.nn.model_parallel import initialize as mpu
|
||||
from transformers import PreTrainedModel
|
||||
|
||||
from models.mixin import PreTrainedModelPeftMixin, return_reference_model, set_reference_model
|
||||
|
||||
|
||||
class PretrainedModelParallelPreSplitMixin(PreTrainedModelPeftMixin):
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
|
||||
*model_args,
|
||||
**kwargs,
|
||||
):
|
||||
if mpu.model_parallel_is_initialized():
|
||||
mp_rank = mpu.get_model_parallel_rank()
|
||||
pretrained_model_name_or_path = os.path.join(pretrained_model_name_or_path, f"mp_{mp_rank}-of-{mpu.get_model_parallel_world_size()}")
|
||||
print(f"Loading model from {pretrained_model_name_or_path}")
|
||||
return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
if mpu.model_parallel_is_initialized():
|
||||
mp_rank = mpu.get_model_parallel_rank()
|
||||
save_directory = os.path.join(save_directory, f"mp_{mp_rank}-of-{mpu.get_model_parallel_world_size()}")
|
||||
super().save_pretrained(save_directory, *args, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained_with_ref_model(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], ref_model: PreTrainedModel,
|
||||
*model_args, **kwargs):
|
||||
set_reference_model(ref_model)
|
||||
ref_model = return_reference_model()
|
||||
ref_model.eval()
|
||||
ref_model.to(device=torch.cuda.current_device())
|
||||
|
||||
model = cls.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
|
||||
return model
|
||||
@@ -0,0 +1,717 @@
|
||||
import os
|
||||
from logging import Logger
|
||||
from typing import Optional, Union, Tuple, List, Callable
|
||||
|
||||
import omegaconf
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers.models.llama.modeling_llama import (
|
||||
LlamaForCausalLM as HfLlamaForCausalLM,
|
||||
CausalLMOutputWithPast,
|
||||
LlamaModel,
|
||||
LlamaPreTrainedModel,
|
||||
LlamaConfig,
|
||||
)
|
||||
|
||||
from general_util.logger import get_child_logger
|
||||
from models.dpo_utils import (
|
||||
llama_dpo_batch_forward,
|
||||
llama_last_token_cls_batch_forward,
|
||||
llama_token_batch_forward,
|
||||
llama_last_token_forward_value,
|
||||
llama_batch_forward,
|
||||
sft_loss_on_logits,
|
||||
tdpo_get_batch_logps,
|
||||
)
|
||||
from models.mixin import PreTrainedModelPeftMixin, return_reference_model
|
||||
from models.utils import DPOModelOutput, RewardModelOutput
|
||||
|
||||
logger: Logger = get_child_logger(__name__)
|
||||
|
||||
|
||||
def return_single_device_map():
|
||||
return {"": "cuda:" + str(int(os.environ.get("LOCAL_RANK") or 0))}
|
||||
|
||||
|
||||
class LlamaForCausalLMDPO(PreTrainedModelPeftMixin, HfLlamaForCausalLM):
|
||||
def __init__(self, config, beta: float = 0.1, label_smoothing: float = 0.0, use_ipo: bool = False, loss_type: str = "sigmoid",
|
||||
sft_loss: bool = False, sft_loss_weight: float = 1.0):
|
||||
super().__init__(config)
|
||||
self.beta = beta
|
||||
self.label_smoothing = label_smoothing
|
||||
self.use_ipo = use_ipo
|
||||
self.loss_type = loss_type
|
||||
self.sft_loss = sft_loss
|
||||
self.sft_loss_weight = sft_loss_weight
|
||||
logger.warning(f"Using loss type: {self.loss_type}")
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
|
||||
def dpo_loss(
|
||||
self,
|
||||
policy_chosen_logps: torch.FloatTensor,
|
||||
policy_rejected_logps: torch.FloatTensor,
|
||||
reference_chosen_logps: torch.FloatTensor,
|
||||
reference_rejected_logps: torch.FloatTensor,
|
||||
reference_free: bool = False,
|
||||
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
"""Compute the DPO loss for a batch of policy and reference model log probabilities.
|
||||
|
||||
Args:
|
||||
policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)
|
||||
policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)
|
||||
reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,)
|
||||
reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,)
|
||||
beta: Temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
|
||||
reference_free: If True, we ignore the _provided_ reference model and implicitly use a reference model that assigns equal probability to all responses.
|
||||
|
||||
Returns:
|
||||
A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).
|
||||
The losses tensor contains the DPO loss for each example in the batch.
|
||||
The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.
|
||||
"""
|
||||
pi_logratios = policy_chosen_logps - policy_rejected_logps
|
||||
ref_logratios = reference_chosen_logps - reference_rejected_logps
|
||||
|
||||
if reference_free:
|
||||
ref_logratios = 0
|
||||
|
||||
logits = pi_logratios - ref_logratios
|
||||
|
||||
if self.use_ipo:
|
||||
losses = (logits - 1 / (2 * self.beta)) ** 2
|
||||
elif self.loss_type == "hinge":
|
||||
losses = torch.relu(1 - self.beta * logits)
|
||||
elif self.loss_type == "sigmoid":
|
||||
log_sigmoid = nn.LogSigmoid()
|
||||
losses = -log_sigmoid(self.beta * logits) * (1 - self.label_smoothing) - log_sigmoid(-self.beta * logits) * self.label_smoothing
|
||||
else:
|
||||
raise ValueError(f"Unsupported loss type: {self.loss_type}")
|
||||
|
||||
chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps).detach()
|
||||
rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps).detach()
|
||||
|
||||
return losses.mean(), chosen_rewards, rejected_rewards
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
policy_logits, policy_logprobs, policy_loss_mask = llama_dpo_batch_forward(self, input_ids, attention_mask, labels)
|
||||
with torch.no_grad():
|
||||
ref_logits, ref_logprobs, ref_loss_mask = llama_dpo_batch_forward(return_reference_model(), input_ids, attention_mask, labels,
|
||||
pad_token_id=self.config.pad_token_id)
|
||||
|
||||
policy_chosen_logits, policy_reject_logits = policy_logits[:half], policy_logits[half:]
|
||||
policy_chosen_logprobs, policy_reject_logprobs = policy_logprobs[:half], policy_logprobs[half:]
|
||||
|
||||
ref_chosen_logprobs, ref_reject_logprobs = ref_logprobs[:half], ref_logprobs[half:]
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.dpo_loss(
|
||||
policy_chosen_logps=policy_chosen_logprobs,
|
||||
policy_rejected_logps=policy_reject_logprobs,
|
||||
reference_chosen_logps=ref_chosen_logprobs,
|
||||
reference_rejected_logps=ref_reject_logprobs,
|
||||
reference_free=False,
|
||||
)
|
||||
|
||||
if self.sft_loss:
|
||||
sft_loss = sft_loss_on_logits(policy_chosen_logits, labels[:half], self.config.pad_token_id)
|
||||
loss += self.sft_loss_weight * sft_loss
|
||||
else:
|
||||
sft_loss = None
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_chosen_logits,
|
||||
policy_rejected_logits=policy_reject_logits,
|
||||
sft_loss=sft_loss,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["LlamaForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to LlamaForCausalLM")
|
||||
|
||||
|
||||
class LlamaForCausalLMKTO(PreTrainedModelPeftMixin, HfLlamaForCausalLM):
|
||||
def __init__(self, config, beta: float = 0.1, desirable_weight: float = 1.0, undesirable_weight: float = 1.0):
|
||||
super().__init__(config)
|
||||
self.beta = beta
|
||||
self.desirable_weight = desirable_weight
|
||||
self.undesirable_weight = undesirable_weight
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
|
||||
def kto_loss(
|
||||
self,
|
||||
policy_logps: torch.FloatTensor,
|
||||
reference_logps: torch.FloatTensor,
|
||||
policy_kl_logps: torch.FloatTensor,
|
||||
reference_kl_logps: torch.FloatTensor,
|
||||
desirable_mask: torch.LongTensor,
|
||||
reference_free: bool = False,
|
||||
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
raise NotImplementedError
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
policy_logits, policy_logprobs, policy_loss_mask = llama_dpo_batch_forward(self, input_ids, attention_mask, labels)
|
||||
with torch.no_grad():
|
||||
ref_logits, ref_logprobs, ref_loss_mask = llama_dpo_batch_forward(return_reference_model(), input_ids, attention_mask, labels)
|
||||
|
||||
policy_chosen_logits, policy_reject_logits = policy_logits[:half], policy_logits[half:]
|
||||
policy_chosen_logprobs, policy_reject_logprobs = policy_logprobs[:half], policy_logprobs[half:]
|
||||
|
||||
# ref_chosen_logits, ref_reject_logits = ref_logits[:half], ref_logits[half:]
|
||||
ref_chosen_logprobs, ref_reject_logprobs = ref_logprobs[:half], ref_logprobs[half:]
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.dpo_loss(
|
||||
policy_chosen_logps=policy_chosen_logprobs,
|
||||
policy_rejected_logps=policy_reject_logprobs,
|
||||
reference_chosen_logps=ref_chosen_logprobs,
|
||||
reference_rejected_logps=ref_reject_logprobs,
|
||||
reference_free=False,
|
||||
)
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_chosen_logits,
|
||||
policy_rejected_logits=policy_reject_logits,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["LlamaForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to LlamaForCausalLM")
|
||||
|
||||
|
||||
class LlamaForCausalLMSimPO(PreTrainedModelPeftMixin, HfLlamaForCausalLM):
|
||||
def __init__(self, config, gamma: float, beta: float = 0.1, label_smoothing: float = 0.0, use_ipo: bool = False, loss_type: str = "sigmoid",
|
||||
sft_loss: bool = False, sft_loss_weight: float = 1.0):
|
||||
super().__init__(config)
|
||||
self.gamma = gamma
|
||||
self.beta = beta
|
||||
self.label_smoothing = label_smoothing
|
||||
self.use_ipo = use_ipo
|
||||
self.loss_type = loss_type
|
||||
self.sft_loss = sft_loss
|
||||
self.sft_loss_weight = sft_loss_weight
|
||||
logger.warning(f"Using loss type: {self.loss_type}")
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
|
||||
def simpo_loss(
|
||||
self,
|
||||
policy_chosen_logps: torch.FloatTensor,
|
||||
policy_rejected_logps: torch.FloatTensor,
|
||||
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
"""Compute the SimPO loss for a batch of policy model log probabilities.
|
||||
|
||||
Args:
|
||||
policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)
|
||||
policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)
|
||||
|
||||
Returns:
|
||||
A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).
|
||||
The losses tensor contains the SimPO loss for each example in the batch.
|
||||
The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.
|
||||
"""
|
||||
pi_logratios = policy_chosen_logps - policy_rejected_logps
|
||||
gamma_logratios = self.gamma / self.beta
|
||||
# pi_logratios = pi_logratios.to(self.accelerator.device)
|
||||
logits = pi_logratios - gamma_logratios
|
||||
|
||||
if self.loss_type == "sigmoid":
|
||||
log_sigmoid = nn.LogSigmoid()
|
||||
losses = -log_sigmoid(self.beta * logits) * (1 - self.label_smoothing) - log_sigmoid(-self.beta * logits) * self.label_smoothing
|
||||
elif self.loss_type == "hinge":
|
||||
losses = torch.relu(1 - self.beta * logits)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown loss type: {self.loss_type}. Should be one of ['sigmoid', 'hinge']"
|
||||
)
|
||||
|
||||
chosen_rewards = self.beta * policy_chosen_logps.detach()
|
||||
rejected_rewards = self.beta * policy_rejected_logps.detach()
|
||||
|
||||
return losses.mean(), chosen_rewards, rejected_rewards
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
policy_logits, policy_logprobs, policy_loss_mask = llama_dpo_batch_forward(self, input_ids, attention_mask, labels, average_log_prob=True)
|
||||
|
||||
policy_chosen_logits, policy_reject_logits = policy_logits[:half], policy_logits[half:]
|
||||
policy_chosen_logprobs, policy_reject_logprobs = policy_logprobs[:half], policy_logprobs[half:]
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.simpo_loss(
|
||||
policy_chosen_logps=policy_chosen_logprobs,
|
||||
policy_rejected_logps=policy_reject_logprobs,
|
||||
)
|
||||
|
||||
if self.sft_loss:
|
||||
sft_loss = sft_loss_on_logits(policy_chosen_logits, labels[:half], self.config.pad_token_id)
|
||||
loss += self.sft_loss_weight * sft_loss
|
||||
else:
|
||||
sft_loss = None
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_chosen_logits,
|
||||
policy_rejected_logits=policy_reject_logits,
|
||||
sft_loss=sft_loss,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["LlamaForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to LlamaForCausalLM")
|
||||
|
||||
|
||||
class LlamaForCausalLMTDPO(PreTrainedModelPeftMixin, HfLlamaForCausalLM):
|
||||
def __init__(self, config, beta: float, alpha: float = 0.5, sft_loss: bool = False, sft_loss_weight: float = 1.0, if_tdpo2: bool = True, ):
|
||||
super().__init__(config)
|
||||
self.beta = beta
|
||||
self.alpha = alpha
|
||||
self.sft_loss = sft_loss
|
||||
self.sft_loss_weight = sft_loss_weight
|
||||
self.if_tdpo2 = if_tdpo2
|
||||
# logger.warning(f"Using loss type: {self.loss_type}")
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
|
||||
def tdpo_loss(self, chosen_logps_margin: torch.FloatTensor,
|
||||
rejected_logps_margin: torch.FloatTensor,
|
||||
chosen_position_kl: torch.FloatTensor,
|
||||
rejected_position_kl: torch.FloatTensor,
|
||||
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
"""Compute the TDPO loss for a batch of policy and reference model log probabilities.
|
||||
|
||||
Args:
|
||||
chosen_logps_margin: The difference of log probabilities between the policy model and the reference model for the chosen responses. Shape: (batch_size,)
|
||||
rejected_logps_margin: The difference of log probabilities between the policy model and the reference model for the rejected responses. Shape: (batch_size,)
|
||||
chosen_position_kl: The difference of sequential kl divergence between the policy model and the reference model for the chosen responses. Shape: (batch_size,)
|
||||
rejected_position_kl: The difference of sequential kl divergence between the policy model and the reference model for the rejected responses. Shape: (batch_size,)
|
||||
beta: Temperature parameter for the TDPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
|
||||
alpha: Temperature parameter for the TDPO loss, used to adjust the impact of sequential kl divergence.
|
||||
if_tdpo2: Determine whether to use method TDPO2, default is True; if False, then use method TDPO1.
|
||||
|
||||
Returns:
|
||||
A tuple of two tensors: (losses, rewards).
|
||||
The losses tensor contains the TDPO loss for each example in the batch.
|
||||
The rewards tensors contain the rewards for response pair.
|
||||
"""
|
||||
|
||||
chosen_values = chosen_logps_margin + chosen_position_kl
|
||||
rejected_values = rejected_logps_margin + rejected_position_kl
|
||||
|
||||
chosen_rejected_logps_margin = chosen_logps_margin - rejected_logps_margin
|
||||
|
||||
if not self.if_tdpo2:
|
||||
logits = chosen_rejected_logps_margin - (rejected_position_kl - chosen_position_kl) # tdpo1
|
||||
else:
|
||||
logits = chosen_rejected_logps_margin - self.alpha * (rejected_position_kl - chosen_position_kl.detach()) # tdpo2
|
||||
|
||||
log_sigmoid = torch.nn.LogSigmoid()
|
||||
losses = -log_sigmoid(self.beta * logits)
|
||||
|
||||
chosen_rewards = self.beta * chosen_values.detach()
|
||||
rejected_rewards = self.beta * rejected_values.detach()
|
||||
|
||||
return losses.mean(), chosen_rewards, rejected_rewards
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
policy_logits = llama_batch_forward(self, input_ids, attention_mask).to(torch.float32)
|
||||
with torch.no_grad():
|
||||
ref_logits = llama_batch_forward(return_reference_model(), input_ids, attention_mask).to(torch.float32)
|
||||
|
||||
logps_margin, position_kl, logps = tdpo_get_batch_logps(policy_logits, ref_logits, labels, self.config.pad_token_id,
|
||||
average_log_prob=False)
|
||||
|
||||
chosen_logps_margin, rejected_logps_margin = logps_margin[:half], logps_margin[half:]
|
||||
chosen_position_kl, rejected_position_kl = position_kl[:half], position_kl[half:]
|
||||
|
||||
chosen_logps, rejected_logps = logps[:half].detach(), logps[half:].detach()
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.tdpo_loss(
|
||||
chosen_logps_margin=chosen_logps_margin,
|
||||
rejected_logps_margin=rejected_logps_margin,
|
||||
chosen_position_kl=chosen_position_kl,
|
||||
rejected_position_kl=rejected_position_kl,
|
||||
)
|
||||
|
||||
if self.sft_loss:
|
||||
sft_loss = sft_loss_on_logits(policy_logits[:half], labels[:half], self.config.pad_token_id)
|
||||
loss += self.sft_loss_weight * sft_loss
|
||||
else:
|
||||
sft_loss = None
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_logits[:half],
|
||||
policy_rejected_logits=policy_logits[half:],
|
||||
sft_loss=sft_loss,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["LlamaForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to LlamaForCausalLM")
|
||||
|
||||
|
||||
class LlamaRewardModel(PreTrainedModelPeftMixin, LlamaPreTrainedModel):
|
||||
def __init__(self, config: LlamaConfig, use_token_avg: bool = False):
|
||||
super().__init__(config)
|
||||
self.model = LlamaModel(config)
|
||||
self.score = nn.Linear(config.hidden_size, 1, bias=False)
|
||||
self.use_token_avg = use_token_avg
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
|
||||
def pair_wise_loss(self,
|
||||
chosen_rewards: torch.FloatTensor,
|
||||
rejected_rewards: torch.FloatTensor, ):
|
||||
reward_loss = -torch.log(torch.sigmoid(chosen_rewards - rejected_rewards)).mean()
|
||||
return reward_loss
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
if self.use_token_avg:
|
||||
rewards = llama_token_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id, average=True)
|
||||
else:
|
||||
rewards, _ = llama_last_token_cls_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
|
||||
chosen_rewards, rejected_rewards = rewards[:half], rewards[half:]
|
||||
|
||||
loss = self.pair_wise_loss(chosen_rewards, rejected_rewards)
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=None,
|
||||
policy_rejected_logits=None,
|
||||
batch_chosen_reward=chosen_rewards,
|
||||
batch_rejected_reward=rejected_rewards,
|
||||
)
|
||||
|
||||
|
||||
class LlamaRewardModelForEval(LlamaRewardModel):
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
if self.use_token_avg:
|
||||
rewards = llama_token_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id, average=True)
|
||||
else:
|
||||
rewards, _ = llama_last_token_cls_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
|
||||
return DPOModelOutput(
|
||||
batch_chosen_reward=rewards,
|
||||
)
|
||||
|
||||
|
||||
class LlamaModelForSequenceClassification(PreTrainedModelPeftMixin, LlamaPreTrainedModel):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__(config)
|
||||
self.model = LlamaModel(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
values: Optional[torch.LongTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
rewards, sequence_lengths = llama_last_token_cls_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
loss = loss_fct(rewards, values)
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
logits=rewards,
|
||||
)
|
||||
|
||||
|
||||
class LlamaModelForSequenceClassificationForEval(LlamaModelForSequenceClassification):
|
||||
def __init__(self, config: LlamaConfig, return_full_logits: bool = True):
|
||||
super().__init__(config)
|
||||
self.model = LlamaModel(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
self.return_full_logits = return_full_logits
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
rewards, sequence_lengths = llama_last_token_cls_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id,
|
||||
return_full_logits=self.return_full_logits)
|
||||
if self.return_full_logits:
|
||||
return DPOModelOutput(
|
||||
logits=rewards,
|
||||
)
|
||||
return DPOModelOutput(
|
||||
batch_chosen_reward=rewards,
|
||||
)
|
||||
|
||||
|
||||
class LlamaModelForSequenceClassificationForRL(LlamaModelForSequenceClassification):
|
||||
def __init__(self, config: LlamaConfig, reduce_func: Callable):
|
||||
super().__init__(config)
|
||||
|
||||
self.reduce_func = reduce_func
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, RewardModelOutput]:
|
||||
values, rewards, sequence_lengths = llama_last_token_forward_value(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
values = self.reduce_func(values)
|
||||
rewards = self.reduce_func(rewards)
|
||||
|
||||
value_mask = input_ids.eq(self.config.pad_token_id)
|
||||
values = values.masked_fill(value_mask, 0)
|
||||
|
||||
return RewardModelOutput(
|
||||
values=values,
|
||||
chosen_end_scores=rewards,
|
||||
sequence_lengths=sequence_lengths,
|
||||
)
|
||||
|
||||
|
||||
class LlamaForCausalLM(PreTrainedModelPeftMixin, HfLlamaForCausalLM):
|
||||
def forward(self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
if self.config.pretraining_tp > 1:
|
||||
lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
|
||||
logits = [nn.functional.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
|
||||
logits = torch.cat(logits, dim=-1)
|
||||
else:
|
||||
logits = self.lm_head(hidden_states)
|
||||
logits = logits.float()
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
# Enable model parallelism
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
shift_labels[shift_labels.eq(self.config.pad_token_id)] = -100
|
||||
loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
@@ -0,0 +1,487 @@
|
||||
import os
|
||||
from typing import Union, Optional, Callable, List, Tuple
|
||||
|
||||
import os
|
||||
from typing import Union, Optional, Callable, List, Tuple
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
from megatron.core import parallel_state as mpu
|
||||
from megatron.core.model_parallel_config import ModelParallelConfig
|
||||
from megatron.core.tensor_parallel.layers import ColumnParallelLinear as ColumnParallelLinearMP, RowParallelLinear as RowParallelLinearMP, \
|
||||
VocabParallelEmbedding
|
||||
from megatron.core.tensor_parallel.utils import VocabUtility
|
||||
from torch import nn
|
||||
from torch.nn import init
|
||||
from transformers.models.llama import modeling_llama
|
||||
from transformers.models.llama.configuration_llama import LlamaConfig
|
||||
from transformers.models.llama.modeling_llama import (
|
||||
LlamaAttention,
|
||||
LlamaFlashAttention2,
|
||||
LlamaSdpaAttention,
|
||||
LlamaMLP,
|
||||
LlamaDecoderLayer,
|
||||
LlamaRMSNorm,
|
||||
LlamaModel,
|
||||
is_flash_attn_greater_or_equal_2_10,
|
||||
LlamaForCausalLM as HfLlamaForCausalLM,
|
||||
LlamaPreTrainedModel,
|
||||
CausalLMOutputWithPast,
|
||||
)
|
||||
|
||||
from general_util.logger import get_child_logger
|
||||
from models.dpo_utils import llama_last_token_forward_value, llama_dpo_batch_forward, sft_loss_on_logits, llama_last_token_cls_batch_forward
|
||||
from models.megatron_tp_mixin import PretrainedModelParallelPreSplitMixin
|
||||
from models.mixin import return_reference_model
|
||||
from models.utils import DPOModelOutput, RewardModelOutput
|
||||
|
||||
logger = get_child_logger(__name__)
|
||||
|
||||
_MEGATRON_MP_CONFIG: ModelParallelConfig
|
||||
|
||||
|
||||
class ColumnParallelLinear(ColumnParallelLinearMP):
|
||||
def forward(self, input_: torch.Tensor, weight: Optional[torch.Tensor] = None):
|
||||
return super().forward(input_, weight)[0]
|
||||
|
||||
|
||||
class RowParallelLinear(RowParallelLinearMP):
|
||||
def forward(self, input_):
|
||||
return super().forward(input_)[0]
|
||||
|
||||
|
||||
def init_megatron_mp_config(*args, **kwargs):
|
||||
config = ModelParallelConfig(
|
||||
*args,
|
||||
tensor_model_parallel_size=mpu.get_tensor_model_parallel_world_size(),
|
||||
pipeline_model_parallel_size=mpu.get_pipeline_model_parallel_world_size(),
|
||||
**kwargs,
|
||||
)
|
||||
global _MEGATRON_MP_CONFIG
|
||||
_MEGATRON_MP_CONFIG = config
|
||||
|
||||
|
||||
def attention_tp_init(self: LlamaAttention, config: LlamaConfig):
|
||||
self.q_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_heads * self.head_dim,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=config.attention_bias,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.k_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_key_value_heads * self.head_dim,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.v_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_key_value_heads * self.head_dim,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.o_proj = RowParallelLinear(
|
||||
self.hidden_size,
|
||||
self.hidden_size,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
init_method=lambda x: x,
|
||||
skip_bias_add=False,
|
||||
)
|
||||
if hasattr(self, "_init_rope"):
|
||||
self._init_rope()
|
||||
|
||||
# self.output_size_per_partition = self.q_proj.output_size_per_partition
|
||||
self.num_heads = self.num_heads // mpu.get_tensor_model_parallel_world_size()
|
||||
self.num_key_value_heads = self.num_key_value_heads // mpu.get_tensor_model_parallel_world_size()
|
||||
self.hidden_size = self.hidden_size // mpu.get_tensor_model_parallel_world_size()
|
||||
|
||||
|
||||
class LlamaAttentionParallel(LlamaAttention):
|
||||
def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
|
||||
class LlamaFlashAttention2Parallel(LlamaFlashAttention2):
|
||||
def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
|
||||
|
||||
|
||||
class LlamaSdpaAttentionParallel(LlamaSdpaAttention):
|
||||
"""
|
||||
Llama attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
|
||||
`LlamaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
|
||||
SDPA API.
|
||||
"""
|
||||
|
||||
# Adapted from LlamaAttention.forward
|
||||
def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
|
||||
class LlamaMLPParallel(LlamaMLP):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__(config)
|
||||
|
||||
self.gate_proj = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
config.intermediate_size,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.up_proj = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
config.intermediate_size,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.down_proj = RowParallelLinear(
|
||||
config.intermediate_size,
|
||||
config.hidden_size,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
init_method=lambda x: x,
|
||||
skip_bias_add=False,
|
||||
)
|
||||
|
||||
|
||||
# modeling_llama.LlamaAttention = LlamaAttentionParallel
|
||||
# modeling_llama.LlamaMLP = LlamaMLPParallel
|
||||
# modeling_llama.LLAMA_ATTENTION_CLASSES["eager"] = LlamaAttentionParallel
|
||||
# modeling_llama.LLAMA_ATTENTION_CLASSES["flash_attention_2"] = LlamaFlashAttention2Parallel
|
||||
# modeling_llama.LLAMA_ATTENTION_CLASSES["sdpa"] = LlamaSdpaAttentionParallel
|
||||
|
||||
|
||||
class LlamaModelParallel(LlamaModel):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__(config)
|
||||
self.padding_idx = config.pad_token_id
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.vocab_start_index, self.vocab_end_index = VocabUtility.vocab_range_from_global_vocab_size(
|
||||
config.vocab_size, mpu.get_tensor_model_parallel_rank(), mpu.get_tensor_model_parallel_world_size()
|
||||
)
|
||||
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
# padding_idx=self.padding_idx if config.pad_token_id != config.eos_token_id else None, # TODO: Not sure if this is correct.
|
||||
# This should be consistent with the non-parallel version.
|
||||
# padding_idx=self.padding_idx - self.vocab_start_index if self.vocab_start_index <= self.padding_idx < self.vocab_end_index else None,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
init_method=init.xavier_normal_,
|
||||
)
|
||||
self.layers = nn.ModuleList(
|
||||
[LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
||||
)
|
||||
self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
# # register a causal mask to separate causal and padding mask creation. Merging happends in the attention class
|
||||
# causal_mask = torch.full((config.max_position_embeddings, config.max_position_embeddings), fill_value=1)
|
||||
# self.register_buffer("causal_mask", torch.triu(causal_mask, diagonal=1), persistent=False)
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
|
||||
class LlamaForCausalLM(PretrainedModelParallelPreSplitMixin, HfLlamaForCausalLM):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__(config)
|
||||
|
||||
orig_attention_class = (modeling_llama.LLAMA_ATTENTION_CLASSES["eager"], modeling_llama.LLAMA_ATTENTION_CLASSES["flash_attention_2"],
|
||||
modeling_llama.LLAMA_ATTENTION_CLASSES["sdpa"])
|
||||
orig_mlp_class = modeling_llama.LlamaMLP
|
||||
|
||||
modeling_llama.LlamaAttention = LlamaAttentionParallel
|
||||
modeling_llama.LlamaMLP = LlamaMLPParallel
|
||||
modeling_llama.LLAMA_ATTENTION_CLASSES["eager"] = LlamaAttentionParallel
|
||||
modeling_llama.LLAMA_ATTENTION_CLASSES["flash_attention_2"] = LlamaFlashAttention2Parallel
|
||||
modeling_llama.LLAMA_ATTENTION_CLASSES["sdpa"] = LlamaSdpaAttentionParallel
|
||||
|
||||
self.model = LlamaModelParallel(config)
|
||||
self.lm_head = ColumnParallelLinear(config.hidden_size, config.vocab_size, bias=False, config=_MEGATRON_MP_CONFIG, init_method=lambda x: x,
|
||||
gather_output=True)
|
||||
|
||||
modeling_llama.LlamaAttention = orig_attention_class[0]
|
||||
modeling_llama.LlamaMLP = orig_mlp_class
|
||||
modeling_llama.LLAMA_ATTENTION_CLASSES["eager"] = orig_attention_class[0]
|
||||
modeling_llama.LLAMA_ATTENTION_CLASSES["flash_attention_2"] = orig_attention_class[1]
|
||||
modeling_llama.LLAMA_ATTENTION_CLASSES["sdpa"] = orig_attention_class[2]
|
||||
|
||||
self.post_init()
|
||||
|
||||
def forward(self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
logits = self.lm_head(hidden_states)
|
||||
logits = logits.float()
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
with torch.cuda.amp.autocast(enabled=True, dtype=torch.float32):
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
# Enable model parallelism
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
shift_labels[shift_labels.eq(self.config.pad_token_id)] = -100 # Take care of here.
|
||||
loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
|
||||
class LlamaModelForSequenceClassification(PretrainedModelParallelPreSplitMixin, LlamaPreTrainedModel):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__(config)
|
||||
self.model = LlamaModelParallel(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
values: Optional[torch.LongTensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
rewards, sequence_lengths = llama_last_token_cls_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
loss = loss_fct(rewards, values)
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
logits=rewards,
|
||||
)
|
||||
|
||||
|
||||
class LlamaModelForSequenceClassificationForRL(PretrainedModelParallelPreSplitMixin, LlamaPreTrainedModel):
|
||||
def __init__(self, config: LlamaConfig, reduce_func: Callable):
|
||||
super().__init__(config)
|
||||
self.model = LlamaModelParallel(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
self.reduce_func = reduce_func
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, RewardModelOutput]:
|
||||
values, rewards, sequence_lengths = llama_last_token_forward_value(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
values = self.reduce_func(values)
|
||||
rewards = self.reduce_func(rewards)
|
||||
|
||||
value_mask = input_ids.eq(self.config.pad_token_id)
|
||||
values = values.masked_fill(value_mask, 0)
|
||||
|
||||
return RewardModelOutput(
|
||||
values=values,
|
||||
chosen_end_scores=rewards,
|
||||
sequence_lengths=sequence_lengths,
|
||||
)
|
||||
|
||||
|
||||
class LlamaForCausalLMDPO(LlamaForCausalLM):
|
||||
def __init__(self, config, beta: float = 0.1, label_smoothing: float = 0.0, use_ipo: bool = False, loss_type: str = "sigmoid",
|
||||
sft_loss: bool = False, sft_loss_weight: float = 1.0):
|
||||
super().__init__(config)
|
||||
self.beta = beta
|
||||
self.label_smoothing = label_smoothing
|
||||
self.use_ipo = use_ipo
|
||||
self.loss_type = loss_type
|
||||
self.sft_loss = sft_loss
|
||||
self.sft_loss_weight = sft_loss_weight
|
||||
logger.warning(f"Using loss type: {self.loss_type}")
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.amp.autocast("cuda", enabled=True, dtype=torch.float32)
|
||||
def dpo_loss(
|
||||
self,
|
||||
policy_chosen_logps: torch.FloatTensor,
|
||||
policy_rejected_logps: torch.FloatTensor,
|
||||
reference_chosen_logps: torch.FloatTensor,
|
||||
reference_rejected_logps: torch.FloatTensor,
|
||||
reference_free: bool = False,
|
||||
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
"""Compute the DPO loss for a batch of policy and reference model log probabilities.
|
||||
|
||||
Args:
|
||||
policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)
|
||||
policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)
|
||||
reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,)
|
||||
reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,)
|
||||
beta: Temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
|
||||
reference_free: If True, we ignore the _provided_ reference model and implicitly use a reference model that assigns equal probability to all responses.
|
||||
|
||||
Returns:
|
||||
A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).
|
||||
The losses tensor contains the DPO loss for each example in the batch.
|
||||
The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.
|
||||
"""
|
||||
pi_logratios = policy_chosen_logps - policy_rejected_logps
|
||||
ref_logratios = reference_chosen_logps - reference_rejected_logps
|
||||
|
||||
if reference_free:
|
||||
ref_logratios = 0
|
||||
|
||||
logits = pi_logratios - ref_logratios
|
||||
|
||||
if self.use_ipo:
|
||||
losses = (logits - 1 / (2 * self.beta)) ** 2
|
||||
elif self.loss_type == "hinge":
|
||||
losses = torch.relu(1 - self.beta * logits)
|
||||
elif self.loss_type == "sigmoid":
|
||||
log_sigmoid = nn.LogSigmoid()
|
||||
losses = -log_sigmoid(self.beta * logits) * (1 - self.label_smoothing) - log_sigmoid(-self.beta * logits) * self.label_smoothing
|
||||
else:
|
||||
raise ValueError(f"Unsupported loss type: {self.loss_type}")
|
||||
|
||||
chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps).detach()
|
||||
rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps).detach()
|
||||
|
||||
return losses.mean(), chosen_rewards, rejected_rewards
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
policy_logits, policy_logprobs, policy_loss_mask = llama_dpo_batch_forward(self, input_ids, attention_mask, labels, self.config.pad_token_id)
|
||||
with torch.no_grad():
|
||||
ref_logits, ref_logprobs, ref_loss_mask = llama_dpo_batch_forward(return_reference_model(), input_ids, attention_mask, labels,
|
||||
self.config.pad_token_id)
|
||||
|
||||
policy_chosen_logits, policy_reject_logits = policy_logits[:half], policy_logits[half:]
|
||||
policy_chosen_logprobs, policy_reject_logprobs = policy_logprobs[:half], policy_logprobs[half:]
|
||||
|
||||
ref_chosen_logprobs, ref_reject_logprobs = ref_logprobs[:half], ref_logprobs[half:]
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.dpo_loss(
|
||||
policy_chosen_logps=policy_chosen_logprobs,
|
||||
policy_rejected_logps=policy_reject_logprobs,
|
||||
reference_chosen_logps=ref_chosen_logprobs,
|
||||
reference_rejected_logps=ref_reject_logprobs,
|
||||
reference_free=False,
|
||||
)
|
||||
|
||||
if self.sft_loss:
|
||||
sft_loss = sft_loss_on_logits(policy_chosen_logits, labels[:half], self.config.pad_token_id)
|
||||
loss += self.sft_loss_weight * sft_loss
|
||||
else:
|
||||
sft_loss = None
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_chosen_logits,
|
||||
policy_rejected_logits=policy_reject_logits,
|
||||
sft_loss=sft_loss,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if mpu.model_parallel_is_initialized():
|
||||
mp_rank = mpu.get_tensor_model_parallel_rank()
|
||||
save_directory = os.path.join(save_directory, f"mp_{mp_rank}-of-{mpu.get_tensor_model_parallel_world_size()}")
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["LlamaForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to LlamaForCausalLM")
|
||||
@@ -0,0 +1,428 @@
|
||||
import os
|
||||
from typing import Union, Optional, Callable, List, Tuple
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
from fairscale.nn.model_parallel import initialize as mpu
|
||||
from fairscale.nn.model_parallel.layers import ColumnParallelLinear, RowParallelLinear, VocabParallelEmbedding
|
||||
from fairscale.nn.model_parallel.utils import VocabUtility
|
||||
from torch import nn
|
||||
from transformers.models.llama import modeling_llama
|
||||
from transformers.models.llama.configuration_llama import LlamaConfig
|
||||
from transformers.models.llama.modeling_llama import (
|
||||
LlamaAttention,
|
||||
LlamaFlashAttention2,
|
||||
LlamaSdpaAttention,
|
||||
LlamaMLP,
|
||||
LlamaDecoderLayer,
|
||||
LlamaRMSNorm,
|
||||
LlamaModel,
|
||||
is_flash_attn_greater_or_equal_2_10,
|
||||
LlamaForCausalLM as HfLlamaForCausalLM,
|
||||
LlamaPreTrainedModel,
|
||||
CausalLMOutputWithPast,
|
||||
)
|
||||
|
||||
from general_util.logger import get_child_logger
|
||||
from models.dpo_utils import llama_last_token_forward_value, llama_dpo_batch_forward, sft_loss_on_logits, llama_last_token_cls_batch_forward
|
||||
from models.fs_tp_mixin import PretrainedModelParallelPreSplitMixin
|
||||
from models.mixin import return_reference_model
|
||||
from models.utils import DPOModelOutput, RewardModelOutput
|
||||
|
||||
logger = get_child_logger(__name__)
|
||||
|
||||
|
||||
def attention_tp_init(self: LlamaAttention, config: LlamaConfig):
|
||||
self.q_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_heads * self.head_dim,
|
||||
bias=config.attention_bias,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.k_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_key_value_heads * self.head_dim,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.v_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_key_value_heads * self.head_dim,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.o_proj = RowParallelLinear(
|
||||
self.hidden_size,
|
||||
self.hidden_size,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
if hasattr(self, "_init_rope"):
|
||||
self._init_rope()
|
||||
|
||||
# self.output_size_per_partition = self.q_proj.output_size_per_partition
|
||||
self.num_heads = self.num_heads // mpu.get_model_parallel_world_size()
|
||||
self.num_key_value_heads = self.num_key_value_heads // mpu.get_model_parallel_world_size()
|
||||
self.hidden_size = self.hidden_size // mpu.get_model_parallel_world_size()
|
||||
|
||||
|
||||
class LlamaAttentionParallel(LlamaAttention):
|
||||
def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
|
||||
class LlamaFlashAttention2Parallel(LlamaFlashAttention2):
|
||||
def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
|
||||
|
||||
|
||||
class LlamaSdpaAttentionParallel(LlamaSdpaAttention):
|
||||
"""
|
||||
Llama attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
|
||||
`LlamaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
|
||||
SDPA API.
|
||||
"""
|
||||
|
||||
# Adapted from LlamaAttention.forward
|
||||
def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
|
||||
class LlamaMLPParallel(LlamaMLP):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__(config)
|
||||
|
||||
self.gate_proj = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
config.intermediate_size,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.up_proj = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
config.intermediate_size,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.down_proj = RowParallelLinear(
|
||||
config.intermediate_size,
|
||||
config.hidden_size,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
|
||||
|
||||
modeling_llama.LlamaAttention = LlamaAttentionParallel
|
||||
modeling_llama.LlamaMLP = LlamaMLPParallel
|
||||
modeling_llama.LLAMA_ATTENTION_CLASSES["eager"] = LlamaAttentionParallel
|
||||
modeling_llama.LLAMA_ATTENTION_CLASSES["flash_attention_2"] = LlamaFlashAttention2Parallel
|
||||
modeling_llama.LLAMA_ATTENTION_CLASSES["sdpa"] = LlamaSdpaAttentionParallel
|
||||
|
||||
|
||||
class LlamaModelParallel(LlamaModel):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__(config)
|
||||
self.padding_idx = config.pad_token_id
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.vocab_start_index, self.vocab_end_index = VocabUtility.vocab_range_from_global_vocab_size(
|
||||
config.vocab_size, mpu.get_model_parallel_rank(), mpu.get_model_parallel_world_size()
|
||||
)
|
||||
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
# padding_idx=self.padding_idx if config.pad_token_id != config.eos_token_id else None, # TODO: Not sure if this is correct.
|
||||
# This should be consistent with the non-parallel version.
|
||||
padding_idx=self.padding_idx - self.vocab_start_index if self.vocab_start_index <= self.padding_idx < self.vocab_end_index else None,
|
||||
)
|
||||
self.layers = nn.ModuleList(
|
||||
[LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
||||
)
|
||||
self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
# # register a causal mask to separate causal and padding mask creation. Merging happends in the attention class
|
||||
# causal_mask = torch.full((config.max_position_embeddings, config.max_position_embeddings), fill_value=1)
|
||||
# self.register_buffer("causal_mask", torch.triu(causal_mask, diagonal=1), persistent=False)
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
|
||||
class LlamaForCausalLM(PretrainedModelParallelPreSplitMixin, HfLlamaForCausalLM):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__(config)
|
||||
self.model = LlamaModelParallel(config)
|
||||
self.lm_head = ColumnParallelLinear(config.hidden_size, config.vocab_size, bias=False)
|
||||
|
||||
self.post_init()
|
||||
|
||||
def forward(self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
logits = self.lm_head(hidden_states)
|
||||
logits = logits.float()
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
with torch.cuda.amp.autocast(enabled=True, dtype=torch.float32):
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
# Enable model parallelism
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
shift_labels[shift_labels.eq(self.config.pad_token_id)] = -100 # Take care of here.
|
||||
loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
|
||||
class LlamaModelForSequenceClassification(PretrainedModelParallelPreSplitMixin, LlamaPreTrainedModel):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__(config)
|
||||
self.model = LlamaModelParallel(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
values: Optional[torch.LongTensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
rewards, sequence_lengths = llama_last_token_cls_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
loss = loss_fct(rewards, values)
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
logits=rewards,
|
||||
)
|
||||
|
||||
|
||||
class LlamaModelForSequenceClassificationForRL(PretrainedModelParallelPreSplitMixin, LlamaPreTrainedModel):
|
||||
def __init__(self, config: LlamaConfig, reduce_func: Callable):
|
||||
super().__init__(config)
|
||||
self.model = LlamaModelParallel(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
self.reduce_func = reduce_func
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, RewardModelOutput]:
|
||||
values, rewards, sequence_lengths = llama_last_token_forward_value(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
values = self.reduce_func(values)
|
||||
rewards = self.reduce_func(rewards)
|
||||
|
||||
value_mask = input_ids.eq(self.config.pad_token_id)
|
||||
values = values.masked_fill(value_mask, 0)
|
||||
|
||||
return RewardModelOutput(
|
||||
values=values,
|
||||
chosen_end_scores=rewards,
|
||||
sequence_lengths=sequence_lengths,
|
||||
)
|
||||
|
||||
|
||||
class LlamaForCausalLMDPO(LlamaForCausalLM):
|
||||
def __init__(self, config, beta: float = 0.1, label_smoothing: float = 0.0, use_ipo: bool = False, loss_type: str = "sigmoid",
|
||||
sft_loss: bool = False, sft_loss_weight: float = 1.0):
|
||||
super().__init__(config)
|
||||
self.beta = beta
|
||||
self.label_smoothing = label_smoothing
|
||||
self.use_ipo = use_ipo
|
||||
self.loss_type = loss_type
|
||||
self.sft_loss = sft_loss
|
||||
self.sft_loss_weight = sft_loss_weight
|
||||
logger.warning(f"Using loss type: {self.loss_type}")
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.amp.autocast("cuda", enabled=True, dtype=torch.float32)
|
||||
def dpo_loss(
|
||||
self,
|
||||
policy_chosen_logps: torch.FloatTensor,
|
||||
policy_rejected_logps: torch.FloatTensor,
|
||||
reference_chosen_logps: torch.FloatTensor,
|
||||
reference_rejected_logps: torch.FloatTensor,
|
||||
reference_free: bool = False,
|
||||
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
"""Compute the DPO loss for a batch of policy and reference model log probabilities.
|
||||
|
||||
Args:
|
||||
policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)
|
||||
policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)
|
||||
reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,)
|
||||
reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,)
|
||||
beta: Temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
|
||||
reference_free: If True, we ignore the _provided_ reference model and implicitly use a reference model that assigns equal probability to all responses.
|
||||
|
||||
Returns:
|
||||
A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).
|
||||
The losses tensor contains the DPO loss for each example in the batch.
|
||||
The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.
|
||||
"""
|
||||
pi_logratios = policy_chosen_logps - policy_rejected_logps
|
||||
ref_logratios = reference_chosen_logps - reference_rejected_logps
|
||||
|
||||
if reference_free:
|
||||
ref_logratios = 0
|
||||
|
||||
logits = pi_logratios - ref_logratios
|
||||
|
||||
if self.use_ipo:
|
||||
losses = (logits - 1 / (2 * self.beta)) ** 2
|
||||
elif self.loss_type == "hinge":
|
||||
losses = torch.relu(1 - self.beta * logits)
|
||||
elif self.loss_type == "sigmoid":
|
||||
log_sigmoid = nn.LogSigmoid()
|
||||
losses = -log_sigmoid(self.beta * logits) * (1 - self.label_smoothing) - log_sigmoid(-self.beta * logits) * self.label_smoothing
|
||||
else:
|
||||
raise ValueError(f"Unsupported loss type: {self.loss_type}")
|
||||
|
||||
chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps).detach()
|
||||
rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps).detach()
|
||||
|
||||
return losses.mean(), chosen_rewards, rejected_rewards
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
policy_logits, policy_logprobs, policy_loss_mask = llama_dpo_batch_forward(self, input_ids, attention_mask, labels, self.config.pad_token_id)
|
||||
with torch.no_grad():
|
||||
ref_logits, ref_logprobs, ref_loss_mask = llama_dpo_batch_forward(return_reference_model(), input_ids, attention_mask, labels,
|
||||
self.config.pad_token_id)
|
||||
|
||||
policy_chosen_logits, policy_reject_logits = policy_logits[:half], policy_logits[half:]
|
||||
policy_chosen_logprobs, policy_reject_logprobs = policy_logprobs[:half], policy_logprobs[half:]
|
||||
|
||||
ref_chosen_logprobs, ref_reject_logprobs = ref_logprobs[:half], ref_logprobs[half:]
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.dpo_loss(
|
||||
policy_chosen_logps=policy_chosen_logprobs,
|
||||
policy_rejected_logps=policy_reject_logprobs,
|
||||
reference_chosen_logps=ref_chosen_logprobs,
|
||||
reference_rejected_logps=ref_reject_logprobs,
|
||||
reference_free=False,
|
||||
)
|
||||
|
||||
if self.sft_loss:
|
||||
sft_loss = sft_loss_on_logits(policy_chosen_logits, labels[:half], self.config.pad_token_id)
|
||||
loss += self.sft_loss_weight * sft_loss
|
||||
else:
|
||||
sft_loss = None
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_chosen_logits,
|
||||
policy_rejected_logits=policy_reject_logits,
|
||||
sft_loss=sft_loss,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if mpu.model_parallel_is_initialized():
|
||||
mp_rank = mpu.get_model_parallel_rank()
|
||||
save_directory = os.path.join(save_directory, f"mp_{mp_rank}-of-{mpu.get_model_parallel_world_size()}")
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["LlamaForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to LlamaForCausalLM")
|
||||
@@ -0,0 +1,55 @@
|
||||
import collections
|
||||
import os
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
from megatron.core import parallel_state as mpu
|
||||
from transformers import PreTrainedModel
|
||||
|
||||
from models.mixin import PreTrainedModelPeftMixin, return_reference_model, set_reference_model
|
||||
|
||||
|
||||
class PretrainedModelParallelPreSplitMixin(PreTrainedModelPeftMixin):
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
|
||||
*model_args,
|
||||
**kwargs,
|
||||
):
|
||||
if mpu.model_parallel_is_initialized():
|
||||
mp_rank = mpu.get_tensor_model_parallel_rank()
|
||||
pretrained_model_name_or_path = os.path.join(pretrained_model_name_or_path, f"mp_{mp_rank}-of-{mpu.get_tensor_model_parallel_world_size()}")
|
||||
print(f"Loading model from {pretrained_model_name_or_path}")
|
||||
return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
if mpu.model_parallel_is_initialized():
|
||||
mp_rank = mpu.get_tensor_model_parallel_rank()
|
||||
save_directory = os.path.join(save_directory, f"mp_{mp_rank}-of-{mpu.get_tensor_model_parallel_world_size()}")
|
||||
super().save_pretrained(save_directory, *args, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained_with_ref_model(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], ref_model: PreTrainedModel,
|
||||
*model_args, **kwargs):
|
||||
set_reference_model(ref_model)
|
||||
ref_model = return_reference_model()
|
||||
ref_model.eval()
|
||||
ref_model.to(device=torch.cuda.current_device())
|
||||
|
||||
model = cls.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
|
||||
return model
|
||||
|
||||
def state_dict(self, *args, **kwargs):
|
||||
state_dict = super().state_dict(*args, **kwargs)
|
||||
no_extra_state_dict = collections.OrderedDict()
|
||||
for k, v in state_dict.items():
|
||||
if "_extra_state" in k:
|
||||
continue
|
||||
no_extra_state_dict[k] = v
|
||||
return no_extra_state_dict
|
||||
@@ -0,0 +1,408 @@
|
||||
from typing import Optional, Union, Tuple, List, Callable
|
||||
import os
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers.models.mistral.modeling_mistral import (
|
||||
MistralForCausalLM as HfMistralForCausalLM,
|
||||
CausalLMOutputWithPast,
|
||||
MistralConfig,
|
||||
MistralModel,
|
||||
MistralPreTrainedModel
|
||||
)
|
||||
|
||||
from general_util.logger import get_child_logger
|
||||
from models.dpo_utils import (
|
||||
llama_dpo_batch_forward,
|
||||
llama_last_token_cls_batch_forward,
|
||||
llama_token_batch_forward,
|
||||
llama_last_token_forward_value,
|
||||
llama_batch_forward,
|
||||
sft_loss_on_logits,
|
||||
tdpo_get_batch_logps,
|
||||
)
|
||||
from models.mixin import PreTrainedModelPeftMixin, return_reference_model
|
||||
from models.utils import DPOModelOutput, RewardModelOutput
|
||||
|
||||
logger = get_child_logger(__name__)
|
||||
|
||||
|
||||
class MistralForCausalLMDPO(PreTrainedModelPeftMixin, HfMistralForCausalLM):
|
||||
def __init__(self, config, beta: float = 0.1, label_smoothing: float = 0.0, use_ipo: bool = False, loss_type: str = "sigmoid",
|
||||
sft_loss: bool = False, sft_loss_weight: float = 1.0):
|
||||
super().__init__(config)
|
||||
self.beta = beta
|
||||
self.label_smoothing = label_smoothing
|
||||
self.use_ipo = use_ipo
|
||||
self.loss_type = loss_type
|
||||
self.sft_loss = sft_loss
|
||||
self.sft_loss_weight = sft_loss_weight
|
||||
logger.warning(f"Using loss type: {self.loss_type}")
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
|
||||
def dpo_loss(
|
||||
self,
|
||||
policy_chosen_logps: torch.FloatTensor,
|
||||
policy_rejected_logps: torch.FloatTensor,
|
||||
reference_chosen_logps: torch.FloatTensor,
|
||||
reference_rejected_logps: torch.FloatTensor,
|
||||
reference_free: bool = False,
|
||||
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
"""Compute the DPO loss for a batch of policy and reference model log probabilities.
|
||||
|
||||
Args:
|
||||
policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)
|
||||
policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)
|
||||
reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,)
|
||||
reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,)
|
||||
beta: Temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
|
||||
reference_free: If True, we ignore the _provided_ reference model and implicitly use a reference model that assigns equal probability to all responses.
|
||||
|
||||
Returns:
|
||||
A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).
|
||||
The losses tensor contains the DPO loss for each example in the batch.
|
||||
The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.
|
||||
"""
|
||||
pi_logratios = policy_chosen_logps - policy_rejected_logps
|
||||
ref_logratios = reference_chosen_logps - reference_rejected_logps
|
||||
|
||||
if reference_free:
|
||||
ref_logratios = 0
|
||||
|
||||
logits = pi_logratios - ref_logratios
|
||||
|
||||
if self.use_ipo:
|
||||
losses = (logits - 1 / (2 * self.beta)) ** 2
|
||||
elif self.loss_type == "hinge":
|
||||
losses = torch.relu(1 - self.beta * logits)
|
||||
elif self.loss_type == "sigmoid":
|
||||
log_sigmoid = nn.LogSigmoid()
|
||||
losses = -log_sigmoid(self.beta * logits) * (1 - self.label_smoothing) - log_sigmoid(-self.beta * logits) * self.label_smoothing
|
||||
else:
|
||||
raise ValueError(f"Unsupported loss type: {self.loss_type}")
|
||||
|
||||
chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps).detach()
|
||||
rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps).detach()
|
||||
|
||||
return losses.mean(), chosen_rewards, rejected_rewards
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
policy_logits, policy_logprobs, policy_loss_mask = llama_dpo_batch_forward(self, input_ids, attention_mask, labels)
|
||||
with torch.no_grad():
|
||||
ref_logits, ref_logprobs, ref_loss_mask = llama_dpo_batch_forward(return_reference_model(), input_ids, attention_mask, labels,
|
||||
pad_token_id=self.config.pad_token_id)
|
||||
|
||||
policy_chosen_logits, policy_reject_logits = policy_logits[:half], policy_logits[half:]
|
||||
policy_chosen_logprobs, policy_reject_logprobs = policy_logprobs[:half], policy_logprobs[half:]
|
||||
|
||||
ref_chosen_logprobs, ref_reject_logprobs = ref_logprobs[:half], ref_logprobs[half:]
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.dpo_loss(
|
||||
policy_chosen_logps=policy_chosen_logprobs,
|
||||
policy_rejected_logps=policy_reject_logprobs,
|
||||
reference_chosen_logps=ref_chosen_logprobs,
|
||||
reference_rejected_logps=ref_reject_logprobs,
|
||||
reference_free=False,
|
||||
)
|
||||
|
||||
if self.sft_loss:
|
||||
sft_loss = sft_loss_on_logits(policy_chosen_logits, labels[:half], self.config.pad_token_id)
|
||||
loss += self.sft_loss_weight * sft_loss
|
||||
else:
|
||||
sft_loss = None
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_chosen_logits,
|
||||
policy_rejected_logits=policy_reject_logits,
|
||||
sft_loss=sft_loss,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["MistralForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to MistralForCausalLM")
|
||||
|
||||
|
||||
class MistralForCausalLMTDPO(PreTrainedModelPeftMixin, HfMistralForCausalLM):
|
||||
def __init__(self, config, beta: float, alpha: float = 0.5, sft_loss: bool = False, sft_loss_weight: float = 1.0, if_tdpo2: bool = True, ):
|
||||
super().__init__(config)
|
||||
self.beta = beta
|
||||
self.alpha = alpha
|
||||
self.sft_loss = sft_loss
|
||||
self.sft_loss_weight = sft_loss_weight
|
||||
self.if_tdpo2 = if_tdpo2
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
|
||||
def tdpo_loss(self, chosen_logps_margin: torch.FloatTensor,
|
||||
rejected_logps_margin: torch.FloatTensor,
|
||||
chosen_position_kl: torch.FloatTensor,
|
||||
rejected_position_kl: torch.FloatTensor,
|
||||
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
"""Compute the TDPO loss for a batch of policy and reference model log probabilities.
|
||||
|
||||
Args:
|
||||
chosen_logps_margin: The difference of log probabilities between the policy model and the reference model for the chosen responses. Shape: (batch_size,)
|
||||
rejected_logps_margin: The difference of log probabilities between the policy model and the reference model for the rejected responses. Shape: (batch_size,)
|
||||
chosen_position_kl: The difference of sequential kl divergence between the policy model and the reference model for the chosen responses. Shape: (batch_size,)
|
||||
rejected_position_kl: The difference of sequential kl divergence between the policy model and the reference model for the rejected responses. Shape: (batch_size,)
|
||||
beta: Temperature parameter for the TDPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
|
||||
alpha: Temperature parameter for the TDPO loss, used to adjust the impact of sequential kl divergence.
|
||||
if_tdpo2: Determine whether to use method TDPO2, default is True; if False, then use method TDPO1.
|
||||
|
||||
Returns:
|
||||
A tuple of two tensors: (losses, rewards).
|
||||
The losses tensor contains the TDPO loss for each example in the batch.
|
||||
The rewards tensors contain the rewards for response pair.
|
||||
"""
|
||||
|
||||
chosen_values = chosen_logps_margin + chosen_position_kl
|
||||
rejected_values = rejected_logps_margin + rejected_position_kl
|
||||
|
||||
chosen_rejected_logps_margin = chosen_logps_margin - rejected_logps_margin
|
||||
|
||||
if not self.if_tdpo2:
|
||||
logits = chosen_rejected_logps_margin - (rejected_position_kl - chosen_position_kl) # tdpo1
|
||||
else:
|
||||
logits = chosen_rejected_logps_margin - self.alpha * (rejected_position_kl - chosen_position_kl.detach()) # tdpo2
|
||||
|
||||
log_sigmoid = torch.nn.LogSigmoid()
|
||||
losses = -log_sigmoid(self.beta * logits)
|
||||
|
||||
chosen_rewards = self.beta * chosen_values.detach()
|
||||
rejected_rewards = self.beta * rejected_values.detach()
|
||||
|
||||
return losses.mean(), chosen_rewards, rejected_rewards
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
policy_logits = llama_batch_forward(self, input_ids, attention_mask).to(torch.float32)
|
||||
with torch.no_grad():
|
||||
ref_logits = llama_batch_forward(return_reference_model(), input_ids, attention_mask).to(torch.float32)
|
||||
|
||||
logps_margin, position_kl, logps = tdpo_get_batch_logps(policy_logits, ref_logits, labels, self.config.pad_token_id,
|
||||
average_log_prob=False)
|
||||
|
||||
chosen_logps_margin, rejected_logps_margin = logps_margin[:half], logps_margin[half:]
|
||||
chosen_position_kl, rejected_position_kl = position_kl[:half], position_kl[half:]
|
||||
|
||||
chosen_logps, rejected_logps = logps[:half].detach(), logps[half:].detach()
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.tdpo_loss(
|
||||
chosen_logps_margin=chosen_logps_margin,
|
||||
rejected_logps_margin=rejected_logps_margin,
|
||||
chosen_position_kl=chosen_position_kl,
|
||||
rejected_position_kl=rejected_position_kl,
|
||||
)
|
||||
|
||||
if self.sft_loss:
|
||||
sft_loss = sft_loss_on_logits(policy_logits[:half], labels[:half], self.config.pad_token_id)
|
||||
loss += self.sft_loss_weight * sft_loss
|
||||
else:
|
||||
sft_loss = None
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_logits[:half],
|
||||
policy_rejected_logits=policy_logits[half:],
|
||||
sft_loss=sft_loss,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["MistralForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to MistralForCausalLM")
|
||||
|
||||
|
||||
class MistralForCausalLM(PreTrainedModelPeftMixin, HfMistralForCausalLM):
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
logits = self.lm_head(hidden_states)
|
||||
logits = logits.float()
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
# Enable model parallelism
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
shift_labels[shift_labels.eq(self.config.pad_token_id)] = -100
|
||||
loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
|
||||
class MistralForSequenceClassification(PreTrainedModelPeftMixin, MistralPreTrainedModel):
|
||||
def __init__(self, config: MistralConfig):
|
||||
super().__init__(config)
|
||||
self.model = MistralModel(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
values: Optional[torch.LongTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
rewards, sequence_lengths = llama_last_token_cls_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
loss = loss_fct(rewards, values)
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
logits=rewards,
|
||||
)
|
||||
|
||||
|
||||
class MistralForSequenceClassificationForEval(MistralForSequenceClassification):
|
||||
def __init__(self, config: MistralConfig, return_full_logits: bool = True):
|
||||
super().__init__(config)
|
||||
self.model = MistralModel(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
self.return_full_logits = return_full_logits
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
rewards, sequence_lengths = llama_last_token_cls_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id,
|
||||
return_full_logits=self.return_full_logits)
|
||||
if self.return_full_logits:
|
||||
return DPOModelOutput(
|
||||
logits=rewards,
|
||||
)
|
||||
return DPOModelOutput(
|
||||
batch_chosen_reward=rewards,
|
||||
)
|
||||
@@ -0,0 +1,481 @@
|
||||
import os
|
||||
from typing import Optional, List, Tuple, Union, Callable
|
||||
|
||||
import torch
|
||||
from fairscale.nn.model_parallel import initialize as mpu
|
||||
from fairscale.nn.model_parallel.layers import ColumnParallelLinear, RowParallelLinear, VocabParallelEmbedding
|
||||
from fairscale.nn.model_parallel.utils import VocabUtility
|
||||
from torch import nn
|
||||
from transformers.models.mistral import modeling_mistral
|
||||
from transformers.models.mistral.modeling_mistral import (
|
||||
MistralForCausalLM as HfMistralForCausalLM,
|
||||
CausalLMOutputWithPast,
|
||||
MistralAttention,
|
||||
MistralFlashAttention2,
|
||||
MistralSdpaAttention,
|
||||
MistralModel,
|
||||
is_flash_attn_greater_or_equal_2_10,
|
||||
MistralMLP,
|
||||
MistralConfig,
|
||||
MISTRAL_ATTENTION_CLASSES,
|
||||
)
|
||||
|
||||
from general_util.logger import get_child_logger
|
||||
from models.dpo_utils import (
|
||||
llama_dpo_batch_forward,
|
||||
llama_batch_forward,
|
||||
sft_loss_on_logits,
|
||||
tdpo_get_batch_logps,
|
||||
)
|
||||
from models.mixin import PretrainedModelParallelPreSplitMixin, return_reference_model
|
||||
from models.utils import DPOModelOutput
|
||||
|
||||
logger = get_child_logger(__name__)
|
||||
|
||||
|
||||
def attention_tp_init(self: MistralAttention, config: MistralConfig):
|
||||
self.q_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_heads * self.head_dim,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.k_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_key_value_heads * self.head_dim,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.v_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_key_value_heads * self.head_dim,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.o_proj = RowParallelLinear(
|
||||
self.hidden_size,
|
||||
self.hidden_size,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
|
||||
self.num_heads = self.num_heads // mpu.get_model_parallel_world_size()
|
||||
self.num_key_value_heads = self.num_key_value_heads // mpu.get_model_parallel_world_size()
|
||||
self.hidden_size = self.hidden_size // mpu.get_model_parallel_world_size()
|
||||
|
||||
|
||||
class MistralAttentionTensorParallel(MistralAttention):
|
||||
def __init__(self, config: MistralConfig, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
|
||||
class MistralFlashAttentionTensorParallel(MistralFlashAttention2):
|
||||
def __init__(self, config: MistralConfig, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
|
||||
|
||||
|
||||
class MistralSdpaAttentionTensorParallel(MistralSdpaAttention):
|
||||
def __init__(self, config: MistralConfig, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
|
||||
class MistralMLPTensorParallel(MistralMLP):
|
||||
def __init__(self, config: MistralConfig):
|
||||
super().__init__(config)
|
||||
|
||||
self.gate_proj = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
config.intermediate_size,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.up_proj = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
config.intermediate_size,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.down_proj = RowParallelLinear(
|
||||
config.intermediate_size,
|
||||
config.hidden_size,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
|
||||
|
||||
MISTRAL_ATTENTION_CLASSES["eager"] = MistralAttentionTensorParallel
|
||||
MISTRAL_ATTENTION_CLASSES["flash_attention_2"] = MistralFlashAttentionTensorParallel
|
||||
MISTRAL_ATTENTION_CLASSES["sdpa"] = MistralSdpaAttentionTensorParallel
|
||||
modeling_mistral.MistralMLP = MistralMLPTensorParallel
|
||||
|
||||
|
||||
class MistralModelTensorParallel(MistralModel):
|
||||
def __init__(self, config: MistralConfig):
|
||||
super().__init__(config)
|
||||
|
||||
self.padding_idx = config.pad_token_id
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.vocab_start_index, self.vocab_end_index = VocabUtility.vocab_range_from_global_vocab_size(
|
||||
config.vocab_size, mpu.get_model_parallel_rank(), mpu.get_model_parallel_world_size()
|
||||
)
|
||||
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
padding_idx=self.padding_idx - self.vocab_start_index if self.vocab_start_index <= self.padding_idx < self.vocab_end_index else None,
|
||||
)
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
|
||||
class MistralForCausalLM(PretrainedModelParallelPreSplitMixin, HfMistralForCausalLM):
|
||||
def __init__(self, config: MistralConfig):
|
||||
super().__init__(config)
|
||||
self.model = MistralModelTensorParallel(config)
|
||||
self.vocab_size = config.vocab_size
|
||||
self.lm_head = ColumnParallelLinear(config.hidden_size, config.vocab_size, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
logits = self.lm_head(hidden_states)
|
||||
logits = logits.float()
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
# Enable model parallelism
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
shift_labels[shift_labels.eq(self.config.pad_token_id)] = -100
|
||||
loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
|
||||
class MistralForCausalLMDPO(MistralForCausalLM):
|
||||
def __init__(self, config, beta: float = 0.1, label_smoothing: float = 0.0, use_ipo: bool = False, loss_type: str = "sigmoid",
|
||||
sft_loss: bool = False, sft_loss_weight: float = 1.0):
|
||||
super().__init__(config)
|
||||
self.beta = beta
|
||||
self.label_smoothing = label_smoothing
|
||||
self.use_ipo = use_ipo
|
||||
self.loss_type = loss_type
|
||||
self.sft_loss = sft_loss
|
||||
self.sft_loss_weight = sft_loss_weight
|
||||
logger.warning(f"Using loss type: {self.loss_type}")
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
|
||||
def dpo_loss(
|
||||
self,
|
||||
policy_chosen_logps: torch.FloatTensor,
|
||||
policy_rejected_logps: torch.FloatTensor,
|
||||
reference_chosen_logps: torch.FloatTensor,
|
||||
reference_rejected_logps: torch.FloatTensor,
|
||||
reference_free: bool = False,
|
||||
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
"""Compute the DPO loss for a batch of policy and reference model log probabilities.
|
||||
|
||||
Args:
|
||||
policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)
|
||||
policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)
|
||||
reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,)
|
||||
reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,)
|
||||
beta: Temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
|
||||
reference_free: If True, we ignore the _provided_ reference model and implicitly use a reference model that assigns equal probability to all responses.
|
||||
|
||||
Returns:
|
||||
A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).
|
||||
The losses tensor contains the DPO loss for each example in the batch.
|
||||
The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.
|
||||
"""
|
||||
pi_logratios = policy_chosen_logps - policy_rejected_logps
|
||||
ref_logratios = reference_chosen_logps - reference_rejected_logps
|
||||
|
||||
if reference_free:
|
||||
ref_logratios = 0
|
||||
|
||||
logits = pi_logratios - ref_logratios
|
||||
|
||||
if self.use_ipo:
|
||||
losses = (logits - 1 / (2 * self.beta)) ** 2
|
||||
elif self.loss_type == "hinge":
|
||||
losses = torch.relu(1 - self.beta * logits)
|
||||
elif self.loss_type == "sigmoid":
|
||||
log_sigmoid = nn.LogSigmoid()
|
||||
losses = -log_sigmoid(self.beta * logits) * (1 - self.label_smoothing) - log_sigmoid(-self.beta * logits) * self.label_smoothing
|
||||
else:
|
||||
raise ValueError(f"Unsupported loss type: {self.loss_type}")
|
||||
|
||||
chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps).detach()
|
||||
rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps).detach()
|
||||
|
||||
return losses.mean(), chosen_rewards, rejected_rewards
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
policy_logits, policy_logprobs, policy_loss_mask = llama_dpo_batch_forward(self, input_ids, attention_mask, labels)
|
||||
with torch.no_grad():
|
||||
ref_logits, ref_logprobs, ref_loss_mask = llama_dpo_batch_forward(return_reference_model(), input_ids, attention_mask, labels,
|
||||
pad_token_id=self.config.pad_token_id)
|
||||
|
||||
policy_chosen_logits, policy_reject_logits = policy_logits[:half], policy_logits[half:]
|
||||
policy_chosen_logprobs, policy_reject_logprobs = policy_logprobs[:half], policy_logprobs[half:]
|
||||
|
||||
ref_chosen_logprobs, ref_reject_logprobs = ref_logprobs[:half], ref_logprobs[half:]
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.dpo_loss(
|
||||
policy_chosen_logps=policy_chosen_logprobs,
|
||||
policy_rejected_logps=policy_reject_logprobs,
|
||||
reference_chosen_logps=ref_chosen_logprobs,
|
||||
reference_rejected_logps=ref_reject_logprobs,
|
||||
reference_free=False,
|
||||
)
|
||||
|
||||
if self.sft_loss:
|
||||
sft_loss = sft_loss_on_logits(policy_chosen_logits, labels[:half], self.config.pad_token_id)
|
||||
loss += self.sft_loss_weight * sft_loss
|
||||
else:
|
||||
sft_loss = None
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_chosen_logits,
|
||||
policy_rejected_logits=policy_reject_logits,
|
||||
sft_loss=sft_loss,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if mpu.model_parallel_is_initialized():
|
||||
mp_rank = mpu.get_model_parallel_rank()
|
||||
save_directory = os.path.join(save_directory, f"mp_{mp_rank}-of-{mpu.get_model_parallel_world_size()}")
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["MistralForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to MistralForCausalLM")
|
||||
|
||||
|
||||
class MistralForCausalLMTDPO(MistralForCausalLM):
|
||||
def __init__(self, config, beta: float, alpha: float = 0.5, sft_loss: bool = False, sft_loss_weight: float = 1.0, if_tdpo2: bool = True, ):
|
||||
super().__init__(config)
|
||||
self.beta = beta
|
||||
self.alpha = alpha
|
||||
self.sft_loss = sft_loss
|
||||
self.sft_loss_weight = sft_loss_weight
|
||||
self.if_tdpo2 = if_tdpo2
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
|
||||
def tdpo_loss(self, chosen_logps_margin: torch.FloatTensor,
|
||||
rejected_logps_margin: torch.FloatTensor,
|
||||
chosen_position_kl: torch.FloatTensor,
|
||||
rejected_position_kl: torch.FloatTensor,
|
||||
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
"""Compute the TDPO loss for a batch of policy and reference model log probabilities.
|
||||
|
||||
Args:
|
||||
chosen_logps_margin: The difference of log probabilities between the policy model and the reference model for the chosen responses. Shape: (batch_size,)
|
||||
rejected_logps_margin: The difference of log probabilities between the policy model and the reference model for the rejected responses. Shape: (batch_size,)
|
||||
chosen_position_kl: The difference of sequential kl divergence between the policy model and the reference model for the chosen responses. Shape: (batch_size,)
|
||||
rejected_position_kl: The difference of sequential kl divergence between the policy model and the reference model for the rejected responses. Shape: (batch_size,)
|
||||
beta: Temperature parameter for the TDPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
|
||||
alpha: Temperature parameter for the TDPO loss, used to adjust the impact of sequential kl divergence.
|
||||
if_tdpo2: Determine whether to use method TDPO2, default is True; if False, then use method TDPO1.
|
||||
|
||||
Returns:
|
||||
A tuple of two tensors: (losses, rewards).
|
||||
The losses tensor contains the TDPO loss for each example in the batch.
|
||||
The rewards tensors contain the rewards for response pair.
|
||||
"""
|
||||
|
||||
chosen_values = chosen_logps_margin + chosen_position_kl
|
||||
rejected_values = rejected_logps_margin + rejected_position_kl
|
||||
|
||||
chosen_rejected_logps_margin = chosen_logps_margin - rejected_logps_margin
|
||||
|
||||
if not self.if_tdpo2:
|
||||
logits = chosen_rejected_logps_margin - (rejected_position_kl - chosen_position_kl) # tdpo1
|
||||
else:
|
||||
logits = chosen_rejected_logps_margin - self.alpha * (rejected_position_kl - chosen_position_kl.detach()) # tdpo2
|
||||
|
||||
log_sigmoid = torch.nn.LogSigmoid()
|
||||
losses = -log_sigmoid(self.beta * logits)
|
||||
|
||||
chosen_rewards = self.beta * chosen_values.detach()
|
||||
rejected_rewards = self.beta * rejected_values.detach()
|
||||
|
||||
return losses.mean(), chosen_rewards, rejected_rewards
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
policy_logits = llama_batch_forward(self, input_ids, attention_mask).to(torch.float32)
|
||||
with torch.no_grad():
|
||||
ref_logits = llama_batch_forward(return_reference_model(), input_ids, attention_mask).to(torch.float32)
|
||||
|
||||
logps_margin, position_kl, logps = tdpo_get_batch_logps(policy_logits, ref_logits, labels, self.config.pad_token_id,
|
||||
average_log_prob=False)
|
||||
|
||||
chosen_logps_margin, rejected_logps_margin = logps_margin[:half], logps_margin[half:]
|
||||
chosen_position_kl, rejected_position_kl = position_kl[:half], position_kl[half:]
|
||||
|
||||
chosen_logps, rejected_logps = logps[:half].detach(), logps[half:].detach()
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.tdpo_loss(
|
||||
chosen_logps_margin=chosen_logps_margin,
|
||||
rejected_logps_margin=rejected_logps_margin,
|
||||
chosen_position_kl=chosen_position_kl,
|
||||
rejected_position_kl=rejected_position_kl,
|
||||
)
|
||||
|
||||
if self.sft_loss:
|
||||
sft_loss = sft_loss_on_logits(policy_logits[:half], labels[:half], self.config.pad_token_id)
|
||||
loss += self.sft_loss_weight * sft_loss
|
||||
else:
|
||||
sft_loss = None
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_logits[:half],
|
||||
policy_rejected_logits=policy_logits[half:],
|
||||
sft_loss=sft_loss,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if mpu.model_parallel_is_initialized():
|
||||
mp_rank = mpu.get_model_parallel_rank()
|
||||
save_directory = os.path.join(save_directory, f"mp_{mp_rank}-of-{mpu.get_model_parallel_world_size()}")
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["MistralForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to MistralForCausalLM")
|
||||
@@ -0,0 +1,120 @@
|
||||
import os
|
||||
from logging import Logger
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
from transformers import PreTrainedModel
|
||||
|
||||
from general_util.logger import get_child_logger
|
||||
|
||||
logger: Logger = get_child_logger(__name__)
|
||||
|
||||
REFERENCE_MODEL: PreTrainedModel
|
||||
|
||||
|
||||
def return_single_device_map():
|
||||
return {"": "cuda:" + str(int(os.environ.get("LOCAL_RANK") or 0))}
|
||||
|
||||
|
||||
def return_reference_model():
|
||||
return REFERENCE_MODEL
|
||||
|
||||
|
||||
def set_reference_model(model: PreTrainedModel):
|
||||
global REFERENCE_MODEL
|
||||
REFERENCE_MODEL = model
|
||||
|
||||
|
||||
class PreTrainedModelPeftMixin(PreTrainedModel):
|
||||
_ref_model: PreTrainedModel = None
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs):
|
||||
gradient_checkpointing = kwargs.pop("gradient_checkpointing", False)
|
||||
model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
|
||||
if gradient_checkpointing:
|
||||
model.config.use_cache = False
|
||||
model.gradient_checkpointing_enable()
|
||||
|
||||
return model
|
||||
|
||||
@classmethod
|
||||
def from_pretrained_with_ref_model(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], ref_model: PreTrainedModel,
|
||||
register_ref_model: bool = False,
|
||||
*model_args, **kwargs):
|
||||
set_reference_model(ref_model)
|
||||
ref_model = return_reference_model()
|
||||
ref_model.eval()
|
||||
ref_model.to(device=torch.cuda.current_device())
|
||||
|
||||
model = cls.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
|
||||
|
||||
if register_ref_model:
|
||||
model._ref_model = ref_model
|
||||
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def deepspeed_set_ref_engine_lazy(ref_model):
|
||||
set_reference_model(ref_model)
|
||||
|
||||
# TODO: This will leads to error when using ROCm since bitsandbytes does not support AMD platform. Consider a lazy import.
|
||||
# @classmethod
|
||||
# def from_pretrained_with_ref_model_lora(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs):
|
||||
# lora_config = kwargs.pop("lora_config", None)
|
||||
# assert lora_config is not None, "lora_config must be provided to enable lora training."
|
||||
#
|
||||
# base_model = cls.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
|
||||
# global REFERENCE_MODEL
|
||||
# REFERENCE_MODEL = base_model
|
||||
#
|
||||
# enable_quantization = "quantization_config" in kwargs
|
||||
#
|
||||
# if lora_config is None:
|
||||
# lora_config = LoraConfig(task_type=TaskType.SEQ_CLS, inference_mode=False, r=8, lora_alpha=32,
|
||||
# lora_dropout=0.1)
|
||||
#
|
||||
# logger.warning(lora_config)
|
||||
# logger.info(lora_config.target_modules.__class__)
|
||||
# if isinstance(lora_config.target_modules, omegaconf.listconfig.ListConfig):
|
||||
# lora_config.target_modules = list(lora_config.target_modules)
|
||||
# elif isinstance(lora_config.target_modules, omegaconf.DictConfig):
|
||||
# lora_config.target_modules = hydra.utils.instantiate(lora_config.target_modules, model=base_model)
|
||||
# else:
|
||||
# raise ValueError(f"Unsupported type of target modules: {lora_config.target_modules.__class__}")
|
||||
#
|
||||
# if isinstance(lora_config.modules_to_save, omegaconf.listconfig.ListConfig):
|
||||
# lora_config.modules_to_save = list(lora_config.modules_to_save)
|
||||
#
|
||||
# logger.info(lora_config.target_modules.__class__)
|
||||
# logger.warning(lora_config.target_modules)
|
||||
#
|
||||
# gradient_checkpointing = base_model.model.gradient_checkpointing
|
||||
# if enable_quantization:
|
||||
# logger.warning(f"Rank {get_rank()} is being loaded with quantization.")
|
||||
# logger.info(f"Quantization config: {kwargs['quantization_config']}")
|
||||
# base_model = prepare_model_for_kbit_training(base_model, use_gradient_checkpointing=gradient_checkpointing)
|
||||
#
|
||||
# model = get_peft_model(base_model, lora_config)
|
||||
#
|
||||
# logger.info(f"Reference model type: {REFERENCE_MODEL.__class__.__name__}")
|
||||
# logger.info(f"Actor model type: {model.__class__.__name__}")
|
||||
#
|
||||
# compute_dtype = kwargs["torch_dtype"]
|
||||
# for name, module in model.named_modules():
|
||||
# if isinstance(module, LoraLayer):
|
||||
# if compute_dtype == torch.bfloat16:
|
||||
# module = module.to(torch.bfloat16)
|
||||
# if 'norm' in name:
|
||||
# module = module.to(torch.float32)
|
||||
# if 'lm_head' in name or 'embed_tokens' in name:
|
||||
# if hasattr(module, 'weight'):
|
||||
# if compute_dtype and module.weight.dtype == torch.float32:
|
||||
# module = module.to(torch.bfloat16)
|
||||
#
|
||||
# model.print_trainable_parameters()
|
||||
#
|
||||
# logger.info(f"Config pad token id after loading pre-trained weights: {model.config.pad_token_id}")
|
||||
# logger.info(model.lm_head.__class__.__name__)
|
||||
#
|
||||
# return model
|
||||
@@ -0,0 +1,318 @@
|
||||
import os
|
||||
from logging import Logger
|
||||
from typing import Optional, Union, Tuple, List, Callable
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers.models.qwen2.modeling_qwen2 import (
|
||||
Qwen2ForCausalLM as HfQwen2ForCausalLM,
|
||||
CausalLMOutputWithPast,
|
||||
Qwen2Config,
|
||||
Qwen2Model
|
||||
)
|
||||
|
||||
from deepspeed.sequence.layer import DistributedAttention
|
||||
from general_util.logger import get_child_logger
|
||||
from models.dpo_utils import llama_dpo_batch_forward, llama_last_token_cls_batch_forward, sft_loss_on_logits
|
||||
from models.mixin import PreTrainedModelPeftMixin, return_reference_model
|
||||
from models.utils import DPOModelOutput
|
||||
|
||||
logger: Logger = get_child_logger(__name__)
|
||||
|
||||
|
||||
class Qwen2ForCausalLM(PreTrainedModelPeftMixin, HfQwen2ForCausalLM):
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
logits = self.lm_head(hidden_states)
|
||||
logits = logits.float()
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
# Enable model parallelism
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
shift_labels[shift_labels.eq(self.config.pad_token_id)] = -100
|
||||
loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
|
||||
class Qwen2RewardModel(PreTrainedModelPeftMixin, HfQwen2ForCausalLM):
|
||||
def __init__(self, config: Qwen2Config):
|
||||
super().__init__(config)
|
||||
self.model = Qwen2Model(config)
|
||||
self.score = nn.Linear(config.hidden_size, 1, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@staticmethod
|
||||
def pair_wise_loss(chosen_rewards: torch.FloatTensor,
|
||||
rejected_rewards: torch.FloatTensor, ):
|
||||
reward_loss = -torch.log(torch.sigmoid(chosen_rewards - rejected_rewards)).mean()
|
||||
return reward_loss
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
rewards, sequence_lengths = llama_last_token_cls_batch_forward(self.model, self.score, input_ids,
|
||||
attention_mask, self.config.pad_token_id)
|
||||
chosen_rewards, rejected_rewards = rewards[:half], rewards[half:]
|
||||
|
||||
loss = self.pair_wise_loss(chosen_rewards, rejected_rewards)
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=None,
|
||||
policy_rejected_logits=None,
|
||||
batch_chosen_reward=chosen_rewards,
|
||||
batch_rejected_reward=rejected_rewards,
|
||||
)
|
||||
|
||||
|
||||
class Qwen2ForSequenceClassification(PreTrainedModelPeftMixin, HfQwen2ForCausalLM):
|
||||
def __init__(self, config: Qwen2Config):
|
||||
super().__init__(config)
|
||||
self.model = Qwen2Model(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
values: Optional[torch.LongTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
rewards, sequence_lengths = llama_last_token_cls_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
loss = loss_fct(rewards, values)
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
logits=rewards,
|
||||
)
|
||||
|
||||
|
||||
class Qwen2ForCausalLMDPO(PreTrainedModelPeftMixin, HfQwen2ForCausalLM):
|
||||
def __init__(self, config, beta: float = 0.1, label_smoothing: float = 0.0, use_ipo: bool = False, loss_type: str = "sigmoid",
|
||||
sft_loss: bool = False, sft_loss_weight: float = 1.0, sft_loss_marco_average: bool = False,
|
||||
double_forward: bool = False):
|
||||
super().__init__(config)
|
||||
self.beta = beta
|
||||
self.label_smoothing = label_smoothing
|
||||
self.use_ipo = use_ipo
|
||||
self.loss_type = loss_type
|
||||
self.sft_loss = sft_loss
|
||||
self.sft_loss_weight = sft_loss_weight
|
||||
self.sft_loss_marco_average = sft_loss_marco_average
|
||||
self.double_forward = double_forward
|
||||
logger.warning(f"Using loss type: {self.loss_type}")
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.cuda.amp.autocast(enabled=True, dtype=torch.float32) # Reduce memory usage.
|
||||
def dpo_loss(
|
||||
self,
|
||||
policy_chosen_logps: torch.FloatTensor,
|
||||
policy_rejected_logps: torch.FloatTensor,
|
||||
reference_chosen_logps: torch.FloatTensor,
|
||||
reference_rejected_logps: torch.FloatTensor,
|
||||
reference_free: bool = False,
|
||||
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
"""Compute the DPO loss for a batch of policy and reference model log probabilities.
|
||||
|
||||
Args:
|
||||
policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)
|
||||
policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)
|
||||
reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,)
|
||||
reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,)
|
||||
beta: Temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
|
||||
reference_free: If True, we ignore the _provided_ reference model and implicitly use a reference model that assigns equal probability to all responses.
|
||||
|
||||
Returns:
|
||||
A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).
|
||||
The losses tensor contains the DPO loss for each example in the batch.
|
||||
The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.
|
||||
"""
|
||||
pi_logratios = policy_chosen_logps - policy_rejected_logps
|
||||
ref_logratios = reference_chosen_logps - reference_rejected_logps
|
||||
|
||||
if reference_free:
|
||||
ref_logratios = 0
|
||||
|
||||
logits = pi_logratios - ref_logratios
|
||||
|
||||
if self.use_ipo:
|
||||
losses = (logits - 1 / (2 * self.beta)) ** 2
|
||||
elif self.loss_type == "hinge":
|
||||
losses = torch.relu(1 - self.beta * logits)
|
||||
elif self.loss_type == "sigmoid":
|
||||
log_sigmoid = nn.LogSigmoid()
|
||||
losses = -log_sigmoid(self.beta * logits) * (1 - self.label_smoothing) - log_sigmoid(-self.beta * logits) * self.label_smoothing
|
||||
else:
|
||||
raise ValueError(f"Unsupported loss type: {self.loss_type}")
|
||||
|
||||
chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps).detach()
|
||||
rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps).detach()
|
||||
|
||||
return losses.mean(), chosen_rewards, rejected_rewards
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
weights: Optional[torch.FloatTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
if self.double_forward:
|
||||
policy_chosen_logits, policy_chosen_logprobs, policy_loss_mask = llama_dpo_batch_forward(self, input_ids[:half], attention_mask[:half],
|
||||
labels[:half])
|
||||
policy_reject_logits, policy_reject_logprobs, _ = llama_dpo_batch_forward(self, input_ids[half:], attention_mask[half:], labels[half:])
|
||||
with torch.no_grad():
|
||||
_, ref_chosen_logprobs, _ = llama_dpo_batch_forward(return_reference_model(), input_ids[:half], attention_mask[:half], labels[:half],
|
||||
pad_token_id=self.config.pad_token_id)
|
||||
_, ref_reject_logprobs, _ = llama_dpo_batch_forward(return_reference_model(), input_ids[half:], attention_mask[half:], labels[half:],
|
||||
pad_token_id=self.config.pad_token_id)
|
||||
|
||||
else:
|
||||
policy_logits, policy_logprobs, policy_loss_mask = llama_dpo_batch_forward(self, input_ids, attention_mask, labels)
|
||||
with torch.no_grad():
|
||||
_, ref_logprobs, _ = llama_dpo_batch_forward(return_reference_model(), input_ids, attention_mask, labels,
|
||||
pad_token_id=self.config.pad_token_id)
|
||||
|
||||
policy_chosen_logits, policy_reject_logits = policy_logits[:half], policy_logits[half:]
|
||||
policy_chosen_logprobs, policy_reject_logprobs = policy_logprobs[:half], policy_logprobs[half:]
|
||||
|
||||
ref_chosen_logprobs, ref_reject_logprobs = ref_logprobs[:half], ref_logprobs[half:]
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.dpo_loss(
|
||||
policy_chosen_logps=policy_chosen_logprobs,
|
||||
policy_rejected_logps=policy_reject_logprobs,
|
||||
reference_chosen_logps=ref_chosen_logprobs,
|
||||
reference_rejected_logps=ref_reject_logprobs,
|
||||
reference_free=False,
|
||||
)
|
||||
|
||||
if self.sft_loss:
|
||||
if weights is not None and self.sft_loss_marco_average is False:
|
||||
logger.warning("Using row weights for SFT loss should enable macro average.")
|
||||
_sft_loss_marco_average = True
|
||||
else:
|
||||
_sft_loss_marco_average = self.sft_loss_marco_average
|
||||
sft_loss = sft_loss_on_logits(policy_chosen_logits, labels[:half], self.config.pad_token_id, macro_average=_sft_loss_marco_average,
|
||||
row_weights=weights)
|
||||
loss += self.sft_loss_weight * sft_loss
|
||||
else:
|
||||
sft_loss = None
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_chosen_logits,
|
||||
policy_rejected_logits=policy_reject_logits,
|
||||
sft_loss=sft_loss,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["Qwen2ForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to Qwen2ForCausalLM")
|
||||
@@ -0,0 +1,471 @@
|
||||
import copy
|
||||
import os
|
||||
from typing import Union, Optional, Callable, List, Tuple
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
from megatron.core import parallel_state as mpu
|
||||
from megatron.core.tensor_parallel.layers import ColumnParallelLinear as ColumnParallelLinearMP, RowParallelLinear as RowParallelLinearMP, \
|
||||
VocabParallelEmbedding
|
||||
from megatron.core.tensor_parallel.cross_entropy import VocabParallelCrossEntropy
|
||||
from megatron.core.tensor_parallel.utils import VocabUtility, divide
|
||||
from megatron.core.models.gpt import GPTModel
|
||||
from torch import nn
|
||||
from torch.nn import init
|
||||
from transformers.models.qwen2 import modeling_qwen2
|
||||
from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
|
||||
from transformers.models.qwen2.modeling_qwen2 import (
|
||||
Qwen2Attention,
|
||||
Qwen2FlashAttention2,
|
||||
Qwen2SdpaAttention,
|
||||
Qwen2MLP,
|
||||
Qwen2DecoderLayer,
|
||||
Qwen2RMSNorm,
|
||||
Qwen2Model,
|
||||
is_flash_attn_greater_or_equal_2_10,
|
||||
Qwen2ForCausalLM as HfQwen2ForCausalLM,
|
||||
Qwen2PreTrainedModel,
|
||||
CausalLMOutputWithPast,
|
||||
)
|
||||
|
||||
from general_util.logger import get_child_logger
|
||||
from models.dpo_utils import llama_last_token_forward_value, llama_dpo_batch_forward, sft_loss_on_logits, llama_last_token_cls_batch_forward
|
||||
from models.mixin import return_reference_model
|
||||
from models.megatron_tp_mixin import PretrainedModelParallelPreSplitMixin
|
||||
from models.utils import DPOModelOutput, RewardModelOutput
|
||||
from megatron.core.model_parallel_config import ModelParallelConfig
|
||||
|
||||
logger = get_child_logger(__name__)
|
||||
|
||||
_MEGATRON_MP_CONFIG: ModelParallelConfig
|
||||
|
||||
|
||||
class ColumnParallelLinear(ColumnParallelLinearMP):
|
||||
def forward(self, input_: torch.Tensor, weight: Optional[torch.Tensor] = None):
|
||||
return super().forward(input_, weight)[0]
|
||||
|
||||
|
||||
class RowParallelLinear(RowParallelLinearMP):
|
||||
def forward(self, input_):
|
||||
return super().forward(input_)[0]
|
||||
|
||||
|
||||
def init_megatron_mp_config(*args, **kwargs):
|
||||
config = ModelParallelConfig(
|
||||
*args,
|
||||
tensor_model_parallel_size=mpu.get_tensor_model_parallel_world_size(),
|
||||
pipeline_model_parallel_size=mpu.get_pipeline_model_parallel_world_size(),
|
||||
**kwargs,
|
||||
)
|
||||
global _MEGATRON_MP_CONFIG
|
||||
_MEGATRON_MP_CONFIG = config
|
||||
|
||||
|
||||
def attention_tp_init(self: Qwen2Attention, config: Qwen2Config):
|
||||
self.q_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_heads * self.head_dim,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=config.attention_bias,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.k_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_key_value_heads * self.head_dim,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.v_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_key_value_heads * self.head_dim,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.o_proj = RowParallelLinear(
|
||||
self.hidden_size,
|
||||
self.hidden_size,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
init_method=lambda x: x,
|
||||
skip_bias_add=False,
|
||||
)
|
||||
if hasattr(self, "_init_rope"):
|
||||
self._init_rope()
|
||||
|
||||
# self.output_size_per_partition = self.q_proj.output_size_per_partition
|
||||
self.num_heads = self.num_heads // mpu.get_tensor_model_parallel_world_size()
|
||||
self.num_key_value_heads = self.num_key_value_heads // mpu.get_tensor_model_parallel_world_size()
|
||||
self.hidden_size = self.hidden_size // mpu.get_tensor_model_parallel_world_size()
|
||||
|
||||
|
||||
class Qwen2AttentionParallel(Qwen2Attention):
|
||||
def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
|
||||
class Qwen2FlashAttention2Parallel(Qwen2FlashAttention2):
|
||||
def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
|
||||
|
||||
|
||||
class Qwen2SdpaAttentionParallel(Qwen2SdpaAttention):
|
||||
"""
|
||||
qwen2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
|
||||
`qwen2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
|
||||
SDPA API.
|
||||
"""
|
||||
|
||||
# Adapted from qwen2Attention.forward
|
||||
def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
|
||||
class Qwen2MLPParallel(Qwen2MLP):
|
||||
def __init__(self, config: Qwen2Config):
|
||||
super().__init__(config)
|
||||
|
||||
self.gate_proj = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
config.intermediate_size,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.up_proj = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
config.intermediate_size,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.down_proj = RowParallelLinear(
|
||||
config.intermediate_size,
|
||||
config.hidden_size,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
init_method=lambda x: x,
|
||||
skip_bias_add=False,
|
||||
)
|
||||
|
||||
|
||||
modeling_qwen2.Qwen2Attention = Qwen2AttentionParallel
|
||||
modeling_qwen2.Qwen2MLP = Qwen2MLPParallel
|
||||
modeling_qwen2.QWEN2_ATTENTION_CLASSES["eager"] = Qwen2AttentionParallel
|
||||
modeling_qwen2.QWEN2_ATTENTION_CLASSES["flash_attention_2"] = Qwen2FlashAttention2Parallel
|
||||
modeling_qwen2.QWEN2_ATTENTION_CLASSES["sdpa"] = Qwen2SdpaAttentionParallel
|
||||
|
||||
|
||||
class Qwen2ModelParallel(Qwen2Model):
|
||||
def __init__(self, config: Qwen2Config):
|
||||
super().__init__(config)
|
||||
self.padding_idx = config.pad_token_id
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.vocab_start_index, self.vocab_end_index = VocabUtility.vocab_range_from_global_vocab_size(
|
||||
config.vocab_size, mpu.get_tensor_model_parallel_rank(), mpu.get_tensor_model_parallel_world_size()
|
||||
)
|
||||
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
# padding_idx=self.padding_idx if config.pad_token_id != config.eos_token_id else None, # TODO: Not sure if this is correct.
|
||||
# This should be consistent with the non-parallel version.
|
||||
# padding_idx=self.padding_idx - self.vocab_start_index if self.vocab_start_index <= self.padding_idx < self.vocab_end_index else None,
|
||||
config=_MEGATRON_MP_CONFIG,
|
||||
init_method=init.xavier_normal_,
|
||||
)
|
||||
self.layers = nn.ModuleList(
|
||||
[Qwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
||||
)
|
||||
self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
# # register a causal mask to separate causal and padding mask creation. Merging happends in the attention class
|
||||
# causal_mask = torch.full((config.max_position_embeddings, config.max_position_embeddings), fill_value=1)
|
||||
# self.register_buffer("causal_mask", torch.triu(causal_mask, diagonal=1), persistent=False)
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
|
||||
class Qwen2ForCausalLM(PretrainedModelParallelPreSplitMixin, HfQwen2ForCausalLM):
|
||||
def __init__(self, config: Qwen2Config):
|
||||
super().__init__(config)
|
||||
|
||||
self.model = Qwen2ModelParallel(config)
|
||||
self.lm_head = ColumnParallelLinear(config.hidden_size, config.vocab_size, bias=False, config=_MEGATRON_MP_CONFIG, init_method=lambda x: x,
|
||||
gather_output=True)
|
||||
|
||||
self.post_init()
|
||||
|
||||
def forward(self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
logits = self.lm_head(hidden_states)
|
||||
logits = logits.float()
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
with torch.cuda.amp.autocast(enabled=True, dtype=torch.float32):
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
# Enable model parallelism
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
shift_labels[shift_labels.eq(self.config.pad_token_id)] = -100 # Take care of here.
|
||||
loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
|
||||
class Qwen2ModelForSequenceClassification(PretrainedModelParallelPreSplitMixin, Qwen2PreTrainedModel):
|
||||
def __init__(self, config: Qwen2Config):
|
||||
super().__init__(config)
|
||||
self.model = Qwen2ModelParallel(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
values: Optional[torch.LongTensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
rewards, sequence_lengths = llama_last_token_cls_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
loss = loss_fct(rewards, values)
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
logits=rewards,
|
||||
)
|
||||
|
||||
|
||||
class Qwen2ModelForSequenceClassificationForRL(PretrainedModelParallelPreSplitMixin, Qwen2PreTrainedModel):
|
||||
def __init__(self, config: Qwen2Config, reduce_func: Callable):
|
||||
super().__init__(config)
|
||||
self.model = Qwen2ModelParallel(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
self.reduce_func = reduce_func
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, RewardModelOutput]:
|
||||
values, rewards, sequence_lengths = llama_last_token_forward_value(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
values = self.reduce_func(values)
|
||||
rewards = self.reduce_func(rewards)
|
||||
|
||||
value_mask = input_ids.eq(self.config.pad_token_id)
|
||||
values = values.masked_fill(value_mask, 0)
|
||||
|
||||
return RewardModelOutput(
|
||||
values=values,
|
||||
chosen_end_scores=rewards,
|
||||
sequence_lengths=sequence_lengths,
|
||||
)
|
||||
|
||||
|
||||
class Qwen2ForCausalLMDPO(Qwen2ForCausalLM):
|
||||
def __init__(self, config, beta: float = 0.1, label_smoothing: float = 0.0, use_ipo: bool = False, loss_type: str = "sigmoid",
|
||||
sft_loss: bool = False, sft_loss_weight: float = 1.0):
|
||||
super().__init__(config)
|
||||
self.beta = beta
|
||||
self.label_smoothing = label_smoothing
|
||||
self.use_ipo = use_ipo
|
||||
self.loss_type = loss_type
|
||||
self.sft_loss = sft_loss
|
||||
self.sft_loss_weight = sft_loss_weight
|
||||
logger.warning(f"Using loss type: {self.loss_type}")
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.amp.autocast("cuda", enabled=True, dtype=torch.float32)
|
||||
def dpo_loss(
|
||||
self,
|
||||
policy_chosen_logps: torch.FloatTensor,
|
||||
policy_rejected_logps: torch.FloatTensor,
|
||||
reference_chosen_logps: torch.FloatTensor,
|
||||
reference_rejected_logps: torch.FloatTensor,
|
||||
reference_free: bool = False,
|
||||
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
"""Compute the DPO loss for a batch of policy and reference model log probabilities.
|
||||
|
||||
Args:
|
||||
policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)
|
||||
policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)
|
||||
reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,)
|
||||
reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,)
|
||||
beta: Temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
|
||||
reference_free: If True, we ignore the _provided_ reference model and implicitly use a reference model that assigns equal probability to all responses.
|
||||
|
||||
Returns:
|
||||
A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).
|
||||
The losses tensor contains the DPO loss for each example in the batch.
|
||||
The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.
|
||||
"""
|
||||
pi_logratios = policy_chosen_logps - policy_rejected_logps
|
||||
ref_logratios = reference_chosen_logps - reference_rejected_logps
|
||||
|
||||
if reference_free:
|
||||
ref_logratios = 0
|
||||
|
||||
logits = pi_logratios - ref_logratios
|
||||
|
||||
if self.use_ipo:
|
||||
losses = (logits - 1 / (2 * self.beta)) ** 2
|
||||
elif self.loss_type == "hinge":
|
||||
losses = torch.relu(1 - self.beta * logits)
|
||||
elif self.loss_type == "sigmoid":
|
||||
log_sigmoid = nn.LogSigmoid()
|
||||
losses = -log_sigmoid(self.beta * logits) * (1 - self.label_smoothing) - log_sigmoid(-self.beta * logits) * self.label_smoothing
|
||||
else:
|
||||
raise ValueError(f"Unsupported loss type: {self.loss_type}")
|
||||
|
||||
chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps).detach()
|
||||
rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps).detach()
|
||||
|
||||
return losses.mean(), chosen_rewards, rejected_rewards
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
policy_logits, policy_logprobs, policy_loss_mask = llama_dpo_batch_forward(self, input_ids, attention_mask, labels, self.config.pad_token_id)
|
||||
with torch.no_grad():
|
||||
ref_logits, ref_logprobs, ref_loss_mask = llama_dpo_batch_forward(return_reference_model(), input_ids, attention_mask, labels,
|
||||
self.config.pad_token_id)
|
||||
|
||||
policy_chosen_logits, policy_reject_logits = policy_logits[:half], policy_logits[half:]
|
||||
policy_chosen_logprobs, policy_reject_logprobs = policy_logprobs[:half], policy_logprobs[half:]
|
||||
|
||||
ref_chosen_logprobs, ref_reject_logprobs = ref_logprobs[:half], ref_logprobs[half:]
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.dpo_loss(
|
||||
policy_chosen_logps=policy_chosen_logprobs,
|
||||
policy_rejected_logps=policy_reject_logprobs,
|
||||
reference_chosen_logps=ref_chosen_logprobs,
|
||||
reference_rejected_logps=ref_reject_logprobs,
|
||||
reference_free=False,
|
||||
)
|
||||
|
||||
if self.sft_loss:
|
||||
sft_loss = sft_loss_on_logits(policy_chosen_logits, labels[:half], self.config.pad_token_id)
|
||||
loss += self.sft_loss_weight * sft_loss
|
||||
else:
|
||||
sft_loss = None
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_chosen_logits,
|
||||
policy_rejected_logits=policy_reject_logits,
|
||||
sft_loss=sft_loss,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if mpu.model_parallel_is_initialized():
|
||||
mp_rank = mpu.get_tensor_model_parallel_rank()
|
||||
save_directory = os.path.join(save_directory, f"mp_{mp_rank}-of-{mpu.get_tensor_model_parallel_world_size()}")
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["qwen2ForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to qwen2ForCausalLM")
|
||||
@@ -0,0 +1,435 @@
|
||||
import os
|
||||
from typing import Union, Optional, Callable, List, Tuple
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
from fairscale.nn.model_parallel import initialize as mpu
|
||||
from fairscale.nn.model_parallel.layers import ColumnParallelLinear, RowParallelLinear, VocabParallelEmbedding
|
||||
from fairscale.nn.model_parallel.utils import VocabUtility
|
||||
from torch import nn
|
||||
from transformers.models.qwen2 import modeling_qwen2
|
||||
from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
|
||||
from transformers.models.qwen2.modeling_qwen2 import (
|
||||
Qwen2Attention,
|
||||
Qwen2FlashAttention2,
|
||||
Qwen2SdpaAttention,
|
||||
Qwen2MLP,
|
||||
Qwen2DecoderLayer,
|
||||
Qwen2RMSNorm,
|
||||
Qwen2Model,
|
||||
is_flash_attn_greater_or_equal_2_10,
|
||||
Qwen2ForCausalLM as HfQwen2ForCausalLM,
|
||||
Qwen2PreTrainedModel,
|
||||
CausalLMOutputWithPast,
|
||||
)
|
||||
|
||||
from general_util.logger import get_child_logger
|
||||
from models.dpo_utils import llama_last_token_forward_value, llama_dpo_batch_forward, sft_loss_on_logits, llama_last_token_cls_batch_forward
|
||||
from models.fs_tp_mixin import PretrainedModelParallelPreSplitMixin
|
||||
from models.mixin import return_reference_model
|
||||
from models.utils import DPOModelOutput, RewardModelOutput
|
||||
|
||||
logger = get_child_logger(__name__)
|
||||
|
||||
|
||||
def attention_tp_init(self: Qwen2Attention, config: Qwen2Config):
|
||||
self.q_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_heads * self.head_dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.k_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_key_value_heads * self.head_dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.v_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.num_key_value_heads * self.head_dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.o_proj = RowParallelLinear(
|
||||
self.hidden_size,
|
||||
self.hidden_size,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
|
||||
# self.output_size_per_partition = self.q_proj.output_size_per_partition
|
||||
self.num_heads = self.num_heads // mpu.get_model_parallel_world_size()
|
||||
self.num_key_value_heads = self.num_key_value_heads // mpu.get_model_parallel_world_size()
|
||||
self.hidden_size = self.hidden_size // mpu.get_model_parallel_world_size()
|
||||
|
||||
|
||||
class Qwen2AttentionParallel(Qwen2Attention):
|
||||
def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
|
||||
class Qwen2FlashAttention2Parallel(Qwen2FlashAttention2):
|
||||
def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
|
||||
|
||||
|
||||
class Qwen2SdpaAttentionParallel(Qwen2SdpaAttention):
|
||||
"""
|
||||
Qwen2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
|
||||
`Qwen2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
|
||||
SDPA API.
|
||||
"""
|
||||
|
||||
# Adapted from Qwen2Attention.forward
|
||||
def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
|
||||
super().__init__(config, layer_idx)
|
||||
|
||||
attention_tp_init(self, config)
|
||||
|
||||
|
||||
class Qwen2MLPParallel(Qwen2MLP):
|
||||
def __init__(self, config: Qwen2Config):
|
||||
super().__init__(config)
|
||||
|
||||
self.gate_proj = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
config.intermediate_size,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.up_proj = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
config.intermediate_size,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
self.down_proj = RowParallelLinear(
|
||||
config.intermediate_size,
|
||||
config.hidden_size,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
init_method=lambda x: x
|
||||
)
|
||||
|
||||
|
||||
modeling_qwen2.Qwen2Attention = Qwen2AttentionParallel
|
||||
modeling_qwen2.Qwen2MLP = Qwen2MLPParallel
|
||||
modeling_qwen2.QWEN2_ATTENTION_CLASSES["eager"] = Qwen2AttentionParallel
|
||||
modeling_qwen2.QWEN2_ATTENTION_CLASSES["flash_attention_2"] = Qwen2FlashAttention2Parallel
|
||||
modeling_qwen2.QWEN2_ATTENTION_CLASSES["sdpa"] = Qwen2SdpaAttentionParallel
|
||||
|
||||
|
||||
class Qwen2ModelParallel(Qwen2Model):
|
||||
def __init__(self, config: Qwen2Config):
|
||||
super().__init__(config)
|
||||
self.padding_idx = config.pad_token_id
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.vocab_start_index, self.vocab_end_index = VocabUtility.vocab_range_from_global_vocab_size(
|
||||
config.vocab_size, mpu.get_model_parallel_rank(), mpu.get_model_parallel_world_size()
|
||||
)
|
||||
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
# padding_idx=self.padding_idx if config.pad_token_id != config.eos_token_id else None, # TODO: Not sure if this is correct.
|
||||
# This should be consistent with the non-parallel version.
|
||||
padding_idx=self.padding_idx - self.vocab_start_index if self.vocab_start_index <= self.padding_idx < self.vocab_end_index else None,
|
||||
)
|
||||
self.layers = nn.ModuleList(
|
||||
[Qwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
||||
)
|
||||
self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
# # register a causal mask to separate causal and padding mask creation. Merging happends in the attention class
|
||||
# causal_mask = torch.full((config.max_position_embeddings, config.max_position_embeddings), fill_value=1)
|
||||
# self.register_buffer("causal_mask", torch.triu(causal_mask, diagonal=1), persistent=False)
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
|
||||
class Qwen2ForCausalLM(PretrainedModelParallelPreSplitMixin, HfQwen2ForCausalLM):
|
||||
def __init__(self, config: Qwen2Config):
|
||||
super().__init__(config)
|
||||
self.model = Qwen2ModelParallel(config)
|
||||
self.lm_head = ColumnParallelLinear(config.hidden_size, config.vocab_size, bias=False)
|
||||
|
||||
self.post_init()
|
||||
|
||||
def forward(self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
logits = self.lm_head(hidden_states)
|
||||
logits = logits.float()
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
with torch.cuda.amp.autocast(enabled=True, dtype=torch.float32):
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
# Enable model parallelism
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
shift_labels[shift_labels.eq(self.config.pad_token_id)] = -100 # Take care of here.
|
||||
loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
|
||||
class Qwen2ModelForSequenceClassification(PretrainedModelParallelPreSplitMixin, Qwen2PreTrainedModel):
|
||||
def __init__(self, config: Qwen2Config):
|
||||
super().__init__(config)
|
||||
self.model = Qwen2ModelParallel(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
values: Optional[torch.LongTensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
rewards, sequence_lengths = llama_last_token_cls_batch_forward(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
loss = loss_fct(rewards, values)
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
logits=rewards,
|
||||
)
|
||||
|
||||
|
||||
class Qwen2ModelForSequenceClassificationForRL(PretrainedModelParallelPreSplitMixin, Qwen2PreTrainedModel):
|
||||
def __init__(self, config: Qwen2Config, reduce_func: Callable):
|
||||
super().__init__(config)
|
||||
self.model = Qwen2ModelParallel(config)
|
||||
self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
|
||||
|
||||
self.reduce_func = reduce_func
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, RewardModelOutput]:
|
||||
values, rewards, sequence_lengths = llama_last_token_forward_value(self.model, self.score, input_ids, attention_mask, self.config.pad_token_id)
|
||||
values = self.reduce_func(values)
|
||||
rewards = self.reduce_func(rewards)
|
||||
|
||||
value_mask = input_ids.eq(self.config.pad_token_id)
|
||||
values = values.masked_fill(value_mask, 0)
|
||||
|
||||
return RewardModelOutput(
|
||||
values=values,
|
||||
chosen_end_scores=rewards,
|
||||
sequence_lengths=sequence_lengths,
|
||||
)
|
||||
|
||||
|
||||
class Qwen2ForCausalLMDPO(Qwen2ForCausalLM):
|
||||
def __init__(self, config, beta: float = 0.1, label_smoothing: float = 0.0, use_ipo: bool = False, loss_type: str = "sigmoid",
|
||||
sft_loss: bool = False, sft_loss_weight: float = 1.0, sft_loss_marco_average: bool = False,):
|
||||
super().__init__(config)
|
||||
self.beta = beta
|
||||
self.label_smoothing = label_smoothing
|
||||
self.use_ipo = use_ipo
|
||||
self.loss_type = loss_type
|
||||
self.sft_loss = sft_loss
|
||||
self.sft_loss_weight = sft_loss_weight
|
||||
self.sft_loss_marco_average = sft_loss_marco_average
|
||||
logger.warning(f"Using loss type: {self.loss_type}")
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
|
||||
def dpo_loss(
|
||||
self,
|
||||
policy_chosen_logps: torch.FloatTensor,
|
||||
policy_rejected_logps: torch.FloatTensor,
|
||||
reference_chosen_logps: torch.FloatTensor,
|
||||
reference_rejected_logps: torch.FloatTensor,
|
||||
reference_free: bool = False,
|
||||
):
|
||||
"""Compute the DPO loss for a batch of policy and reference model log probabilities.
|
||||
|
||||
Args:
|
||||
policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)
|
||||
policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)
|
||||
reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,)
|
||||
reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,)
|
||||
beta: Temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
|
||||
reference_free: If True, we ignore the _provided_ reference model and implicitly use a reference model that assigns equal probability to all responses.
|
||||
|
||||
Returns:
|
||||
A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).
|
||||
The losses tensor contains the DPO loss for each example in the batch.
|
||||
The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.
|
||||
"""
|
||||
pi_logratios = policy_chosen_logps - policy_rejected_logps
|
||||
ref_logratios = reference_chosen_logps - reference_rejected_logps
|
||||
|
||||
if reference_free:
|
||||
ref_logratios = 0
|
||||
|
||||
logits = pi_logratios - ref_logratios
|
||||
|
||||
if self.use_ipo:
|
||||
losses = (logits - 1 / (2 * self.beta)) ** 2
|
||||
elif self.loss_type == "hinge":
|
||||
losses = torch.relu(1 - self.beta * logits)
|
||||
elif self.loss_type == "sigmoid":
|
||||
log_sigmoid = nn.LogSigmoid()
|
||||
losses = -log_sigmoid(self.beta * logits) * (1 - self.label_smoothing) - log_sigmoid(-self.beta * logits) * self.label_smoothing
|
||||
else:
|
||||
raise ValueError(f"Unsupported loss type: {self.loss_type}")
|
||||
|
||||
chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps).detach()
|
||||
rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps).detach()
|
||||
|
||||
return losses.mean(), chosen_rewards, rejected_rewards
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
weights: Optional[torch.FloatTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, DPOModelOutput]:
|
||||
half = input_ids.size(0) // 2
|
||||
|
||||
policy_logits, policy_logprobs, policy_loss_mask = llama_dpo_batch_forward(self, input_ids, attention_mask, labels, self.config.pad_token_id)
|
||||
with torch.no_grad():
|
||||
ref_logits, ref_logprobs, ref_loss_mask = llama_dpo_batch_forward(return_reference_model(), input_ids, attention_mask, labels,
|
||||
self.config.pad_token_id)
|
||||
|
||||
policy_chosen_logits, policy_reject_logits = policy_logits[:half], policy_logits[half:]
|
||||
policy_chosen_logprobs, policy_reject_logprobs = policy_logprobs[:half], policy_logprobs[half:]
|
||||
|
||||
ref_chosen_logprobs, ref_reject_logprobs = ref_logprobs[:half], ref_logprobs[half:]
|
||||
|
||||
loss, chosen_rewards, rejected_rewards = self.dpo_loss(
|
||||
policy_chosen_logps=policy_chosen_logprobs,
|
||||
policy_rejected_logps=policy_reject_logprobs,
|
||||
reference_chosen_logps=ref_chosen_logprobs,
|
||||
reference_rejected_logps=ref_reject_logprobs,
|
||||
reference_free=False,
|
||||
)
|
||||
|
||||
if self.sft_loss:
|
||||
if weights is not None and self.sft_loss_marco_average is False:
|
||||
logger.warning("Using row weights for SFT loss should enable macro average.")
|
||||
_sft_loss_marco_average = True
|
||||
else:
|
||||
_sft_loss_marco_average = self.sft_loss_marco_average
|
||||
sft_loss = sft_loss_on_logits(policy_chosen_logits, labels[:half], self.config.pad_token_id, macro_average=_sft_loss_marco_average,
|
||||
row_weights=weights)
|
||||
loss += self.sft_loss_weight * sft_loss
|
||||
else:
|
||||
sft_loss = None
|
||||
|
||||
return DPOModelOutput(
|
||||
loss=loss,
|
||||
chosen_reward=chosen_rewards.mean(),
|
||||
rejected_reward=rejected_rewards.mean(),
|
||||
policy_chosen_logits=policy_chosen_logits,
|
||||
policy_rejected_logits=policy_reject_logits,
|
||||
sft_loss=sft_loss,
|
||||
)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
state_dict: Optional[dict] = None,
|
||||
save_function: Callable = torch.save,
|
||||
push_to_hub: bool = False,
|
||||
max_shard_size: Union[int, str] = "5GB",
|
||||
safe_serialization: bool = True,
|
||||
variant: Optional[str] = None,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
save_peft_format: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().save_pretrained(save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token,
|
||||
**kwargs)
|
||||
|
||||
if mpu.model_parallel_is_initialized():
|
||||
mp_rank = mpu.get_model_parallel_rank()
|
||||
save_directory = os.path.join(save_directory, f"mp_{mp_rank}-of-{mpu.get_model_parallel_world_size()}")
|
||||
|
||||
if is_main_process:
|
||||
config = self.config
|
||||
config.architectures = ["Qwen2ForCausalLM"]
|
||||
config.save_pretrained(save_directory)
|
||||
logger.warning("Config architecture is override to Qwen2ForCausalLM")
|
||||
@@ -0,0 +1,160 @@
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import hydra.utils
|
||||
import omegaconf
|
||||
import torch
|
||||
from omegaconf import DictConfig
|
||||
from peft import (
|
||||
LoraConfig,
|
||||
get_peft_model,
|
||||
TaskType,
|
||||
prepare_model_for_kbit_training,
|
||||
)
|
||||
from peft.tuners.lora import LoraLayer
|
||||
from transformers import PreTrainedModel
|
||||
from transformers.modeling_outputs import ModelOutput
|
||||
|
||||
from general_util.logger import get_child_logger
|
||||
from general_util.training_utils import get_rank
|
||||
|
||||
logger = get_child_logger(__name__)
|
||||
|
||||
LORA_TARGET_MODULES = [
|
||||
"q_proj",
|
||||
"v_proj",
|
||||
]
|
||||
|
||||
try:
|
||||
import bitsandbytes as bnb
|
||||
except:
|
||||
bnb = None
|
||||
|
||||
|
||||
def find_all_linear_names(model, bits: int, add_lm_head: bool = False):
|
||||
cls = bnb.nn.Linear4bit if bits == 4 else (bnb.nn.Linear8bitLt if bits == 8 else torch.nn.Linear)
|
||||
lora_module_names = set()
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, cls):
|
||||
names = name.split('.')
|
||||
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
|
||||
|
||||
lora_module_names.add("lm_head")
|
||||
|
||||
if 'lm_head' in lora_module_names and not add_lm_head: # needed for 16-bit
|
||||
lora_module_names.remove('lm_head')
|
||||
return list(lora_module_names)
|
||||
|
||||
|
||||
def initialize_peft_model(model: PreTrainedModel, lora_config: DictConfig, load_in_8bit: bool = False, load_in_4bit: bool = False,
|
||||
torch_dtype: torch.dtype = torch.bfloat16):
|
||||
if lora_config is None:
|
||||
lora_config = LoraConfig(task_type=TaskType.CAUSAL_LM, inference_mode=False, r=8, lora_alpha=32,
|
||||
lora_dropout=0.1)
|
||||
|
||||
logger.warning(lora_config)
|
||||
logger.info(lora_config.target_modules.__class__)
|
||||
if isinstance(lora_config.target_modules, omegaconf.listconfig.ListConfig):
|
||||
lora_config.target_modules = list(lora_config.target_modules)
|
||||
elif isinstance(lora_config.target_modules, omegaconf.DictConfig):
|
||||
lora_config.target_modules = hydra.utils.instantiate(lora_config.target_modules, model=model)
|
||||
else:
|
||||
raise ValueError(f"Unsupported type of target modules: {lora_config.target_modules.__class__}")
|
||||
|
||||
if isinstance(lora_config.modules_to_save, omegaconf.listconfig.ListConfig):
|
||||
lora_config.modules_to_save = list(lora_config.modules_to_save)
|
||||
|
||||
logger.warning(lora_config.target_modules)
|
||||
gradient_checkpointing = model.model.gradient_checkpointing
|
||||
if load_in_8bit or load_in_4bit:
|
||||
logger.warning(f"Rank {get_rank()} is being loaded in 8-{load_in_8bit} | 4-{load_in_4bit} bit.")
|
||||
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=gradient_checkpointing)
|
||||
|
||||
model = get_peft_model(model, lora_config)
|
||||
|
||||
compute_dtype = torch_dtype
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, LoraLayer):
|
||||
if compute_dtype == torch.bfloat16:
|
||||
module = module.to(torch.bfloat16)
|
||||
if 'norm' in name:
|
||||
module = module.to(torch.float32)
|
||||
if 'lm_head' in name or 'embed_tokens' in name:
|
||||
if hasattr(module, 'weight'):
|
||||
if compute_dtype and module.weight.dtype == torch.float32:
|
||||
module = module.to(torch.bfloat16)
|
||||
|
||||
model.print_trainable_parameters()
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def enable_gradient_checkpointing(model: PreTrainedModel):
|
||||
model.config.use_cache = False
|
||||
model.gradient_checkpointing_enable()
|
||||
return model
|
||||
|
||||
|
||||
@dataclass
|
||||
class DPOModelOutput(ModelOutput):
|
||||
loss: torch.FloatTensor = None
|
||||
logits: torch.FloatTensor = None
|
||||
chosen_reward: torch.FloatTensor = None
|
||||
rejected_reward: torch.FloatTensor = None
|
||||
policy_chosen_logits: Optional[torch.FloatTensor] = None
|
||||
policy_rejected_logits: Optional[torch.FloatTensor] = None
|
||||
batch_chosen_reward: Optional[torch.FloatTensor] = None
|
||||
batch_rejected_reward: Optional[torch.FloatTensor] = None
|
||||
sft_loss: Optional[torch.FloatTensor] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RewardModelOutput(ModelOutput):
|
||||
values: torch.FloatTensor = None
|
||||
chosen_end_scores: torch.FloatTensor = None
|
||||
sequence_lengths: torch.LongTensor = None
|
||||
|
||||
|
||||
def return_single_device_map():
|
||||
return {"": "cuda:" + str(int(os.environ.get("LOCAL_RANK") or 0))}
|
||||
|
||||
|
||||
def reward_logit2prob(reduction_ids):
|
||||
if isinstance(reduction_ids, omegaconf.ListConfig):
|
||||
reduction_ids = list(reduction_ids)
|
||||
|
||||
def func(logits):
|
||||
probs = torch.softmax(logits, dim=-1)
|
||||
if len(logits.size()) == 3:
|
||||
probs = probs[:, :, reduction_ids].sum(dim=-1)
|
||||
elif len(logits.size()) == 2:
|
||||
probs = probs[:, reduction_ids].sum(dim=-1)
|
||||
else:
|
||||
raise ValueError(f"Unsupported logits shape: {logits.size()}")
|
||||
return probs
|
||||
|
||||
return func
|
||||
|
||||
|
||||
def reward_logit(reduction_ids):
|
||||
if isinstance(reduction_ids, omegaconf.ListConfig):
|
||||
reduction_ids = list(reduction_ids)
|
||||
|
||||
def func(logits):
|
||||
if len(logits.size()) == 3:
|
||||
logits = logits[:, :, reduction_ids].sum(dim=-1)
|
||||
elif len(logits.size()) == 2:
|
||||
logits = logits[:, reduction_ids].sum(dim=-1)
|
||||
else:
|
||||
raise ValueError(f"Unsupported logits shape: {logits.size()}")
|
||||
return logits
|
||||
|
||||
return func
|
||||
|
||||
|
||||
def squeeze_reduce_return_fn():
|
||||
def func(logits: torch.Tensor):
|
||||
return logits.squeeze(-1)
|
||||
|
||||
return func
|
||||
Reference in New Issue
Block a user