chore: import upstream snapshot with attribution

PFPO/models/utils.py

@@ -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