wehub-resource-sync
121 lines
5.0 KiB
Python
121 lines
5.0 KiB
Python
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
|