237 lines
9.1 KiB
Python
237 lines
9.1 KiB
Python
import glob
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Literal, Optional
|
|
|
|
import safetensors
|
|
import torch
|
|
from transformers import TrainingArguments
|
|
|
|
|
|
########## DataClass For Configure ##########
|
|
|
|
|
|
@dataclass
|
|
class TrainingConfig(TrainingArguments):
|
|
max_length: Optional[int] = None
|
|
dataset_num_proc: Optional[int] = None
|
|
center_rewards_coefficient: Optional[float] = None
|
|
disable_flash_attn2: bool = field(default=False)
|
|
|
|
vision_lr: Optional[float] = None
|
|
merger_lr: Optional[float] = None
|
|
special_token_lr: Optional[float] = None
|
|
|
|
conduct_eval: Optional[bool] = True
|
|
load_from_pretrained: str = None
|
|
load_from_pretrained_step: int = None
|
|
logging_epochs: Optional[float] = None
|
|
eval_epochs: Optional[float] = None
|
|
save_epochs: Optional[float] = None
|
|
remove_unused_columns: Optional[bool] = False
|
|
|
|
save_full_model: Optional[bool] = False
|
|
|
|
|
|
@dataclass
|
|
class PEFTLoraConfig:
|
|
lora_enable: bool = False
|
|
vision_lora: bool = False
|
|
lora_r: int = 16
|
|
lora_alpha: int = 32
|
|
lora_dropout: float = 0.05
|
|
lora_target_modules: Optional[List[str]] = None
|
|
lora_namespan_exclude: Optional[List[str]] = None
|
|
lora_modules_to_save: Optional[List[str]] = None
|
|
lora_task_type: str = "CAUSAL_LM"
|
|
use_rslora: bool = False
|
|
num_lora_modules: int = -1
|
|
|
|
def __post_init__(self):
|
|
if isinstance(self.lora_target_modules, list) and len(self.lora_target_modules) == 1:
|
|
self.lora_target_modules = self.lora_target_modules[0]
|
|
|
|
if isinstance(self.lora_namespan_exclude, list) and len(self.lora_namespan_exclude) == 1:
|
|
self.lora_namespan_exclude = self.lora_namespan_exclude[0]
|
|
|
|
|
|
@dataclass
|
|
class ModelConfig:
|
|
model_name_or_path: Optional[str] = None
|
|
model_revision: str = "main"
|
|
|
|
output_dim: int = 1
|
|
|
|
use_special_tokens: bool = False
|
|
|
|
freeze_vision_tower: bool = field(default=False)
|
|
freeze_llm: bool = field(default=False)
|
|
tune_merger: bool = field(default=False)
|
|
|
|
torch_dtype: Optional[Literal["auto", "bfloat16", "float16", "float32"]] = None
|
|
trust_remote_code: bool = False
|
|
attn_implementation: Optional[str] = None
|
|
load_in_8bit: bool = False
|
|
load_in_4bit: bool = False
|
|
bnb_4bit_quant_type: Literal["fp4", "nf4"] = "nf4"
|
|
use_bnb_nested_quant: bool = False
|
|
reward_token: Literal["last", "mean", "special"] = "last"
|
|
loss_type: Literal["bt", "reg", "btt", "margin", "constant_margin", "scaled"] = "regular"
|
|
|
|
def __post_init__(self):
|
|
if self.load_in_8bit and self.load_in_4bit:
|
|
raise ValueError("You can't use 8 bit and 4 bit precision at the same time")
|
|
|
|
# if isinstance(self.lora_target_modules, list) and len(self.lora_target_modules) == 1:
|
|
# self.lora_target_modules = self.lora_target_modules[0]
|
|
|
|
# if isinstance(self.lora_namespan_exclude, list) and len(self.lora_namespan_exclude) == 1:
|
|
# self.lora_namespan_exclude = self.lora_namespan_exclude[0]
|
|
|
|
|
|
########## Functions for get trainable modules' parameters ##########
|
|
|
|
|
|
def maybe_zero_3(param, ignore_status=False, name=None):
|
|
from deepspeed import zero
|
|
|
|
if hasattr(param, "ds_id"):
|
|
# if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
|
|
# if not ignore_status:
|
|
# logging.warning(f"{name}: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: {param.ds_status}")
|
|
with zero.GatheredParameters([param]):
|
|
param = param.data.detach().cpu().clone()
|
|
else:
|
|
param = param.detach().cpu().clone()
|
|
return param
|
|
|
|
|
|
# Borrowed from peft.utils.get_peft_model_state_dict
|
|
def get_peft_state_maybe_zero_3(named_params, bias):
|
|
if bias == "none":
|
|
to_return = {k: t for k, t in named_params if "lora_" in k}
|
|
elif bias == "all":
|
|
to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k}
|
|
elif bias == "lora_only":
|
|
to_return = {}
|
|
maybe_lora_bias = {}
|
|
lora_bias_names = set()
|
|
for k, t in named_params:
|
|
if "lora_" in k:
|
|
to_return[k] = t
|
|
bias_name = k.split("lora_")[0] + "bias"
|
|
lora_bias_names.add(bias_name)
|
|
elif "bias" in k:
|
|
maybe_lora_bias[k] = t
|
|
for k, t in maybe_lora_bias:
|
|
if bias_name in lora_bias_names:
|
|
to_return[bias_name] = t
|
|
else:
|
|
raise NotImplementedError
|
|
to_return = {k: maybe_zero_3(v, ignore_status=True) for k, v in to_return.items()}
|
|
return to_return
|
|
|
|
|
|
def get_peft_state_non_lora_maybe_zero_3(named_params, require_grad_only=True):
|
|
to_return = {k: t for k, t in named_params if "lora_" not in k}
|
|
if require_grad_only:
|
|
to_return = {k: t for k, t in to_return.items() if t.requires_grad}
|
|
to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}
|
|
return to_return
|
|
|
|
|
|
########## Load Models From Folder ##########
|
|
|
|
|
|
def _insert_adapter_name_into_state_dict(
|
|
state_dict: dict[str, torch.Tensor], adapter_name: str, parameter_prefix: str
|
|
) -> dict[str, torch.Tensor]:
|
|
"""Utility function to remap the state_dict keys to fit the PEFT model by inserting the adapter name."""
|
|
peft_model_state_dict = {}
|
|
for key, val in state_dict.items():
|
|
if parameter_prefix in key:
|
|
suffix = key.split(parameter_prefix)[1]
|
|
if "." in suffix:
|
|
suffix_to_replace = ".".join(suffix.split(".")[1:])
|
|
key = key.replace(suffix_to_replace, f"{adapter_name}.{suffix_to_replace}")
|
|
else:
|
|
key = f"{key}.{adapter_name}"
|
|
peft_model_state_dict[key] = val
|
|
else:
|
|
peft_model_state_dict[key] = val
|
|
return peft_model_state_dict
|
|
|
|
|
|
def save_video(tensor, path):
|
|
from torchvision.io import write_video
|
|
|
|
tensor = tensor * 255.0
|
|
tensor = tensor.permute(0, 2, 3, 1)
|
|
tensor = tensor.clamp(0, 255).byte()
|
|
write_video(path, tensor, 4, video_codec="h264")
|
|
|
|
|
|
def load_model_from_checkpoint(model, checkpoint_dir, checkpoint_step):
|
|
checkpoint_paths = glob.glob(os.path.join(checkpoint_dir, "checkpoint-*"))
|
|
checkpoint_paths.sort(key=lambda x: int(x.split("-")[-1]), reverse=True)
|
|
|
|
if checkpoint_step is None or checkpoint_step == -1:
|
|
# get the latest checkpoint
|
|
checkpoint_path = checkpoint_paths[0]
|
|
print(f"===> Checkpoint step is not provided, using the latest checkpoint: {checkpoint_path}")
|
|
else:
|
|
checkpoint_path = os.path.join(checkpoint_dir, f"checkpoint-{checkpoint_step}")
|
|
if checkpoint_path not in checkpoint_paths:
|
|
checkpoint_path = checkpoint_paths[0]
|
|
print(f"===> Checkpoint step {checkpoint_step} not found, using the latest checkpoint: {checkpoint_path}")
|
|
else:
|
|
print(f"===> Checkpoint step {checkpoint_step} found, using the specified checkpoint: {checkpoint_path}")
|
|
|
|
checkpoint_step = checkpoint_path.split("checkpoint-")[-1].split("/")[0]
|
|
|
|
full_ckpt = os.path.join(checkpoint_path, "model.pth")
|
|
lora_ckpt = os.path.join(checkpoint_path, "adapter_model.safetensors")
|
|
non_lora_ckpt = os.path.join(checkpoint_path, "non_lora_state_dict.pth")
|
|
if os.path.exists(full_ckpt):
|
|
model_state_dict = torch.load(full_ckpt, map_location="cpu", weights_only=True)
|
|
# Create a new state_dict to store the modified key-value pairs
|
|
new_state_dict = {}
|
|
|
|
# for key, value in model_state_dict.items():
|
|
# if key.startswith("base_model.model.model"):
|
|
# new_key = "base_model.model.model.language_model" + key[len("base_model.model.model"):]
|
|
# new_state_dict[new_key] = value
|
|
# elif key.startswith("base_model.model.visual"):
|
|
# new_key = "base_model.model.model.visual" + key[len("base_model.model.visual"):]
|
|
# new_state_dict[new_key] = value
|
|
# else:
|
|
# new_state_dict[key] = value
|
|
for key, value in model_state_dict.items():
|
|
if key.startswith("base_model.model.model"):
|
|
new_key = "base_model.model.model.language_model" + key[len("base_model.model.model") :]
|
|
new_state_dict[new_key] = value
|
|
elif key.startswith("base_model.model.visual"):
|
|
new_key = "base_model.model.model.visual" + key[len("base_model.model.visual") :]
|
|
new_state_dict[new_key] = value
|
|
else:
|
|
new_state_dict[key] = value
|
|
|
|
# Load the modified state_dict into the model
|
|
model.load_state_dict(new_state_dict)
|
|
# model_state_dict = torch.load(full_ckpt, map_location="cpu")
|
|
# model.load_state_dict(model_state_dict)
|
|
else:
|
|
lora_state_dict = safetensors.torch.load_file(lora_ckpt)
|
|
non_lora_state_dict = torch.load(non_lora_ckpt, map_location="cpu")
|
|
|
|
lora_state_dict = _insert_adapter_name_into_state_dict(
|
|
lora_state_dict, adapter_name="default", parameter_prefix="lora_"
|
|
)
|
|
|
|
model_state_dict = model.state_dict()
|
|
model_state_dict.update(non_lora_state_dict)
|
|
model_state_dict.update(lora_state_dict)
|
|
model.load_state_dict(model_state_dict)
|
|
|
|
return model, checkpoint_step
|