dpo finetuning
# -*- coding: utf-8 -*-
# @Time:  22:06
# @Author: tk
# @File：__init__.py

from config.main import *



config/__init__.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Time:  21:55
# @Author: tk
# @File：colossalai_config


colossalai_strategy = {
    "ddp": dict(
        name="ddp",
        broadcast_buffers= True,
        bucket_cap_mb = 25,
        find_unused_parameters = False,
        check_reduction = False,
        gradient_as_bucket_view = False,
        static_graph = False,
    ),
    "gemini":dict(
        name="gemini",
        chunk_config_dict = None,
        chunk_init_device= None,
        placement_policy = "static",
        shard_param_frac = 1.0,  # only for static placement
        offload_optim_frac = 0.0,  # only for static placement
        offload_param_frac = 0.0,  # only for static placement
        warmup_non_model_data_ratio = 0.8,  # only for auto placement
        steady_cuda_cap_ratio = 0.9,  # only for auto placement
        precision = "fp16",
        pin_memory = False,
        force_outputs_fp32 = False,
        strict_ddp_mode = False,
        search_range_m = 32,
        hidden_dim = None,
        min_chunk_size_m = 32,
        memstats = None,
        gpu_margin_mem_ratio = 0.0,
        initial_scale = 2 ** 16,
        min_scale = 1,
        growth_factor = 2,
        backoff_factor = 0.5,
        growth_interval = 1000,
        hysteresis = 2,
        max_scale = 2 ** 32,
        max_norm = 1.0,
        norm_type = 2.0,
        verbose = False,
    ),
    "zero2" : dict(
        name="zero2",
        stage = 2,
        precision = "fp16",
        initial_scale = 2 ** 32,
        min_scale = 1,
        growth_factor = 2,
        backoff_factor = 0.5,
        growth_interval = 1000,
        hysteresis = 2,
        max_scale = 2 ** 32,
        max_norm = 1.0,
        norm_type = 2.0,
        reduce_bucket_size_in_m = 12,
        communication_dtype= None,
        overlap_communication = True,
        cpu_offload = False,
        verbose = False,
    ),
    "zero2_cpu" : dict(
        name="zero2_cpu",
        stage = 2,
        precision = "fp16",
        initial_scale = 2 ** 32,
        min_scale = 1,
        growth_factor = 2,
        backoff_factor = 0.5,
        growth_interval = 1000,
        hysteresis = 2,
        max_scale = 2 ** 32,
        max_norm = 1.0,
        norm_type = 2.0,
        reduce_bucket_size_in_m = 12,
        communication_dtype= None,
        overlap_communication = True,
        cpu_offload=True,
        verbose = False,
    ),
    "3d": dict(
        name="3d",
        tp_size =1,
        pp_size = 1,
        precision = "fp16",
        zero_stage = 0,
        enable_all_optimization = False,
        enable_fused_normalization = False,
        enable_flash_attention = False,
        enable_jit_fused = False,
        enable_sequence_parallelism = False,
        enable_sequence_overlap = False,
        num_microbatches = None,
        microbatch_size = None,
        initial_scale = 2 ** 16,
        min_scale = 1,
        growth_factor = 2,
        backoff_factor = 0.5,
        growth_interval = 1000,
        hysteresis = 2,
        max_scale = 2 ** 32,
        max_norm = 0,
        broadcast_buffers = True,
        ddp_bucket_cap_mb = 25,
        find_unused_parameters = False,
        check_reduction = False,
        gradient_as_bucket_view = False,
        static_graph = False,
        zero_bucket_size_in_m = 12,
        cpu_offload = False,
        communication_dtype= None,
        overlap_communication = True,
        custom_policy = None,
    )
}
config/colossalai_config.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Time:  23:20
# @Author: tk
# @File：model_maps
from aigc_zoo.constants.define import (TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING,
                                       TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING,
                                       TRANSFORMERS_MODELS_TO_IA3_TARGET_MODULES_MAPPING,
                                       TRANSFORMERS_MODELS_TO_IA3_FEEDFORWARD_MODULES_MAPPING)

__all__ = [
    "TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING",
    "TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING",
    "TRANSFORMERS_MODELS_TO_IA3_TARGET_MODULES_MAPPING",
    "TRANSFORMERS_MODELS_TO_IA3_FEEDFORWARD_MODULES_MAPPING",
    "MODELS_MAP"
]


MODELS_MAP = {
    'bloom-560m': {
        'model_type': 'bloom',
        'model_name_or_path': '/data/nlp/pre_models/torch/bloom/bloom-560m',
    },
    'bloom-1b7': {
        'model_type': 'bloom',
        'model_name_or_path': '/data/nlp/pre_models/torch/bloom/bloom-1b7',
    },
    'opt-125m': {
        'model_type': 'opt',
        'model_name_or_path': '/data/nlp/pre_models/torch/opt/opt-125m',
    },

    'opt-350m': {
        'model_type': 'opt',
        'model_name_or_path': '/data/nlp/pre_models/torch/opt/opt-350m',
    },

    'llama-7b-hf': {
        'model_type': 'llama',
        'model_name_or_path': '/data/nlp/pre_models/torch/llama/llama-7b-hf',
    },

    'Qwen-7B-Chat': {
        'model_type': 'qwen',
        'model_name_or_path': '/data/nlp/pre_models/torch/qwen/Qwen-7B-Chat',
    },

    'Baichuan2-7B-Chat': {
        'model_type': 'baichuan',
        'model_name_or_path': '/data/nlp/pre_models/torch/baichuan2/Baichuan2-7B-Chat',
    },

    'Baichuan2-13B-Chat': {
        'model_type': 'baichuan',
        'model_name_or_path': '/data/nlp/pre_models/torch/baichuan2/Baichuan2-13B-Chat',
    },

    'Baichuan-13B-Chat': {
        'model_type': 'baichuan',
        'model_name_or_path': '/data/nlp/pre_models/torch/baichuan/Baichuan-13B-Chat',
    },

    'XVERSE-13B-Chat': {
        'model_type': 'xverse',
        'model_name_or_path': '/data/nlp/pre_models/torch/xverse/XVERSE-13B-Chat',
    },

    'internlm-chat-7b': {
        'model_type': 'internlm',
        'model_name_or_path': '/data/nlp/pre_models/torch/internlm/internlm-chat-7b',
    },
    'internlm-chat-7b-8k': {
        'model_type': 'internlm',
        'model_name_or_path': '/data/nlp/pre_models/torch/internlm/internlm-chat-7b-8k',
    },
}


# 按需修改
# TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING
# TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING
# TRANSFORMERS_MODELS_TO_IA3_TARGET_MODULES_MAPPING
# TRANSFORMERS_MODELS_TO_IA3_FEEDFORWARD_MODULES_MAPPING

train_model_config = train_info_models['opt-350m']
config/constant_map.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/5/31 14:43
import json
import os
import torch
import yaml
from transformers import BitsAndBytesConfig
from transformers.utils import strtobool

from config.colossalai_config import colossalai_strategy
from config.constant_map import *



# 加载

__CUR_PATH__ = os.path.abspath(os.path.dirname(__file__))
TRAIN_CONFIG_CACHE = os.path.join(__CUR_PATH__,"config_path_cache.json")


def load_config():
    if "train_config" not in os.environ:
        if not os.path.exists(TRAIN_CONFIG_CACHE):
            raise ValueError("train_config is not set environ")
        with open(TRAIN_CONFIG_CACHE, mode='r', encoding='utf-8') as __f:
            __config_path = json.load(__f)
    else:
        __config_path = os.path.abspath(os.environ.get("train_config"))
        #缓存，方便推理不传环境变量
        with open(TRAIN_CONFIG_CACHE, mode='w', encoding='utf-8') as __f:
            json.dump(__config_path, __f)

    with open(__config_path,mode='r',encoding='utf-8') as __f:
        train_info_args = yaml.full_load(__f)
    return train_info_args

train_info_args = load_config()
global_args = train_info_args.pop("global_args")
train_model_config = MODELS_MAP[global_args["model_name"]]
assert global_args["trainer_backend"] in ["pl","hf","cl","ac"]

# ensure str
global_args["precision"] = str(global_args["precision"])

if global_args["quantization_config"]:
    #精度
    if global_args["precision"] == "auto":
        global_args["quantization_config"]["bnb_4bit_compute_dtype"] ="bfloat16" if torch.cuda.is_bf16_supported() else "float16"

    global_args["quantization_config"] = BitsAndBytesConfig(**global_args["quantization_config"])



def merge_from_env(global_args):
    merge_config = {}
    if "trainer_backend" in os.environ:
        merge_config["trainer_backend"] = str(os.environ["trainer_backend"])
    if "enable_deepspeed" in os.environ:
        merge_config["enable_deepspeed"] = strtobool(os.environ["enable_deepspeed"])
    if "enable_ptv2" in os.environ:
        merge_config["enable_ptv2"] = strtobool(os.environ["enable_ptv2"])
    if "enable_lora" in os.environ:
        merge_config["enable_lora"] = strtobool(os.environ["enable_lora"])
    if "load_in_bit" in os.environ:
        merge_config["load_in_bit"] = int(os.environ["load_in_bit"])
    if merge_config:
        global_args.update(merge_config)

merge_from_env(global_args)

def patch_args(train_info_args):
    assert global_args["enable_lora"] + global_args["enable_ptv2"] <= 1 , ValueError("lora ptv2 cannot open at same time")

    #更新模型配置
    train_info_args.update(train_model_config)

    if global_args["trainer_backend"] == "cl":
        train_info_args["strategy"] = colossalai_strategy[train_info_args["strategy"]]

    if global_args['quantization_config'] is not None:
        global_args['quantization_config'].load_in_4bit = global_args["load_in_bit"] == 4
        global_args['quantization_config'].load_in_8bit = global_args["load_in_bit"] == 8
        if global_args["load_in_bit"] == 0:
            global_args["quantization_config"] = None



    if global_args["enable_lora"]:
        # 检查lora adalora是否开启
        assert train_info_args.get('lora', {}).get('with_lora', False) + \
               train_info_args.get('adalora', {}).get('with_lora', False) + \
               train_info_args.get('ia3', {}).get('with_lora', False) == 1, ValueError(
            'lora adalora ia3 can set one at same time !')

        model_type = train_model_config['model_type']
        if train_info_args.get('lora', {}).get('with_lora', False):
            train_info_args["lora"]["target_modules"] = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_type]
        elif train_info_args.get('adalora', {}).get('with_lora', False):
            train_info_args["adalora"]["target_modules"] = TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING[model_type]
        else:
            train_info_args["ia3"]["target_modules"] = TRANSFORMERS_MODELS_TO_IA3_TARGET_MODULES_MAPPING[model_type]
            train_info_args["ia3"]["feedforward_modules"] = TRANSFORMERS_MODELS_TO_IA3_FEEDFORWARD_MODULES_MAPPING[model_type]

        train_info_args.pop('prompt', None)

    elif global_args["enable_ptv2"]:
        train_info_args.pop('lora', None)
        train_info_args.pop('adalora', None)
        train_info_args.pop('ia3', None)
        if "gradient_checkpointing" in train_info_args:
            train_info_args[ "gradient_checkpointing" ] = False

        assert "prompt" in train_info_args
        train_info_args["prompt"]["with_prompt"] = True
    else:
        train_info_args.pop('lora',None)
        train_info_args.pop('adalora', None)
        train_info_args.pop('ia3', None)
        train_info_args.pop('prompt', None)

    # 预处理
    if 'rwkv' in train_info_args['model_type']:
        train_info_args['use_fast_tokenizer'] = True



patch_args(train_info_args)

def get_deepspeed_config(precision='fp16'):
    '''
        lora prompt finetuning   deepspeed_offload.json
        普通  finetuning          deepspeed.json
    '''
    # 是否开启deepspeed
    if not global_args["enable_deepspeed"]:
        return None
    precision = str(precision).lower()
    # 选择 deepspeed 配置文件
    is_need_update_config = False
    if global_args["enable_lora"] or global_args["enable_ptv2"]:
        is_need_update_config = True
        filename = os.path.join(os.path.dirname(__file__), 'deepspeed_offload.json')
    else:
        filename = os.path.join(os.path.dirname(__file__), 'deepspeed.json')


    with open(filename, mode='r', encoding='utf-8') as f:
        deepspeed_config = json.loads(f.read())

    #lora offload 同步优化器配置
    if is_need_update_config:
        optimizer = deepspeed_config.get('optimizer',None)
        if optimizer:
            if global_args["trainer_backend"] == 'hf':
                optimizer[ 'params' ][ 'betas' ] = (train_info_args.get('adam_beta1', 0.9),train_info_args.get('adam_beta2', 0.999),)
                optimizer[ 'params' ][ 'lr' ] = train_info_args.get('learning_rate', 2e-5)
                optimizer[ 'params' ][ 'eps' ] = train_info_args.get('adam_epsilon', 1e-8)
                # deepspeed_offload 优化器有效
                train_info_args[ 'optim' ] = optimizer[ 'type' ]
            else:
                optimizer['params']['betas'] = train_info_args.get('optimizer_betas', (0.9, 0.999))
                optimizer['params']['lr'] = train_info_args.get('learning_rate', 2e-5)
                optimizer['params']['eps'] = train_info_args.get('adam_epsilon', 1e-8)
                # deepspeed_offload 优化器有效
                train_info_args['optimizer'] = optimizer['type']

    if precision == 'bf16':
        if 'fp16' in deepspeed_config:
            deepspeed_config["fp16"]["enbale"] = False
        if 'bf16' in deepspeed_config:
            deepspeed_config["bf16"]["enbale"] = True
        else:
            deepspeed_config['bf16'] = {"enbale": True}
    elif precision == 'fp16':
        if 'bf16' in deepspeed_config:
            deepspeed_config["bf16"]["enbale"] = False

    return deepspeed_config


config/main.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Time    : 2023/4/20 9:10

from datasets import load_dataset
import json



def make_json_data():

    ds = load_dataset("openai/summarize_from_feedback", name="comparisons")

    # ds['train'].to_json('./train.json')
    # ds['validation'].to_json('./eval.json')

    with open('./train.json', mode='w', encoding='utf-8', newline='\n') as f:
        for d in ds['train']:
            f.write(json.dumps(d, ensure_ascii=False) + '\n')

    with open('./eval.json', mode='w', encoding='utf-8', newline='\n') as f:
        for d in ds['validation']:
            f.write(json.dumps(d, ensure_ascii=False) + '\n')


def make_json_data2():
    ds = load_dataset("Dahoas/rm-static")

    with open('./train.json', mode='w', encoding='utf-8', newline='\n') as f:
        for d in ds['train']:
            f.write(json.dumps(d, ensure_ascii=False) + '\n')

    with open('./eval.json', mode='w', encoding='utf-8', newline='\n') as f:
        for d in ds['test']:
            f.write(json.dumps(d, ensure_ascii=False) + '\n')

def make_json_data3():
    ds = load_dataset("zwh9029/rm-static-m2m100-zh-jianti")

    with open('./train.json', mode='w', encoding='utf-8', newline='\n') as f:
        for d in ds['train']:
            f.write(json.dumps(d, ensure_ascii=False) + '\n')

    with open('./eval.json', mode='w', encoding='utf-8', newline='\n') as f:
        for d in ds['test']:
            f.write(json.dumps(d, ensure_ascii=False) + '\n')

if __name__ == '__main__':
    make_json_data2()
data/make_data_example.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:03

data_factory/__init__.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_baichuan(NN_DataHelper_Base):...
data_factory/data_helper_baichuan.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:03
import copy
import json
import os
import random
import typing
import numpy as np
import torch
from deep_training.data_helper import DataHelper
from aigc_zoo.model_zoo.auto.dpo_model import PetlArguments,LoraConfig
from fastdatasets.record import load_dataset as Loader, RECORD, WriterObject, gfile
from transformers import PreTrainedTokenizer, HfArgumentParser, PretrainedConfig
from data_factory.data_processer import DEFAULT_EOS_TOKEN, DEFAULT_BOS_TOKEN, DEFAULT_UNK_TOKEN, CorpusPreprocess, TokenIdsMaker
from config import *
from torch.nn import functional as F


data_conf = {
    "src_max_length": None,
    "dst_max_length": None,
}


def preprocess(text):
  return text

def postprocess(text):
  return text


class NN_DataHelper_Base(DataHelper):
    index = 1

    def __init__(self, *args, **kwargs):
        super(NN_DataHelper_Base, self).__init__(*args, **kwargs)

    def load_tokenizer_and_config(self, *args, **kwargs):
        ret = super().load_tokenizer_and_config(*args, **kwargs)
        self._preprocess_tokenizer_config()
        return ret
    def _preprocess_tokenizer_config(self):
        model_args = self.model_args
        tokenizer = self.tokenizer
        config = self.config
        if "llama" in model_args.model_name_or_path.lower() and tokenizer.bos_token_id != DEFAULT_BOS_TOKEN:
            tokenizer.add_special_tokens({
                "eos_token": DEFAULT_EOS_TOKEN,
                "bos_token": DEFAULT_BOS_TOKEN,
                "unk_token": DEFAULT_UNK_TOKEN,
            })
            if tokenizer.pad_token_id is None or tokenizer.pad_token_id == -1:
                tokenizer.pad_token_id = tokenizer.eos_token_id

        if tokenizer.pad_token is None:
            tokenizer.add_special_tokens({
                "pad_token": tokenizer.eos_token,
            })
        if config.decoder_start_token_id is None:
            config.decoder_start_token_id = config.bos_token_id
        assert config.decoder_start_token_id == config.bos_token_id



    def on_get_labels(self, files: typing.List[str]):
        D = ['score']
        label2id = {label: i for i, label in enumerate(D)}
        id2label = {i: label for i, label in enumerate(D)}
        return label2id, id2label


    def on_data_ready(self):
        self.index = -1

    # 切分词
    def on_data_process(self, data: typing.Any, mode: str):
        self.index += 1

        tokenizer: PreTrainedTokenizer
        config = self.config
        max_seq_length = self.max_seq_length_dict[mode]
        tokenizer = self.tokenizer

        pair_data = data

        data_conf["sptoken"] = [config.bos_token_id]
        d = TokenIdsMaker.process(pair_data, tokenizer, max_seq_length,**data_conf)
        if self.index < 3:
            print(d)
        return d

    # 读取文件
    def on_get_corpus(self, files: typing.List, mode: str):
        tokenizer = self.tokenizer
        D = []
        for file in files:
            with open(file, mode='r', encoding='utf-8', newline='\n') as f:
                lines = f.readlines()
            d = CorpusPreprocess.process(tokenizer,lines)
            D.extend(d)
        return D

    def collate_fn(self, batch):
        o = {k: [] for k in batch[0].keys()}
        for i, b in enumerate(batch):
            for k in b:
                o[k].append(torch.tensor(b[k]))
        seqlen = np.max([len(_) for _ in o['input_ids']])
        if 'input_ids2' in o:
            seqlen = np.max([seqlen] + [len(_) for _ in o['input_ids2']])

        tokenizer: PreTrainedTokenizer = self.tokenizer
        for k,v in o.items():
            pad_val = tokenizer.pad_token_id if 'label' not in k else -100
            o[k] = torch.stack(
                [F.pad(_, (0, seqlen - len(_)), mode='constant', value=pad_val) for _ in v])
        return o


    def make_dataset_all(self):
        data_args = self.data_args

        # schema for arrow parquet
        schema = None
        # 缓存数据集
        if data_args.do_train:
            self.make_dataset_with_args(data_args.train_file, mixed_data=False, shuffle=True, mode='train',
                                        schema=schema)
        if data_args.do_eval:
            self.make_dataset_with_args(data_args.eval_file, mode='eval', schema=schema)
        if data_args.do_test:
            self.make_dataset_with_args(data_args.test_file, mode='test', schema=schema)

data_factory/data_helper_base.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_bloom(NN_DataHelper_Base):...
data_factory/data_helper_bloom.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from typing import Any
import numpy as np
import torch
from aigc_zoo.model_zoo.chatglm.generation_utils import build_masks_and_position_ids_glm
from torch.nn import functional as F
from transformers import PreTrainedTokenizer
from data_factory.data_helper_base import NN_DataHelper_Base, data_conf
from data_factory.data_processer import TokenIdsMakerForGLM


class NN_DataHelper_chatglm(NN_DataHelper_Base):
    def on_data_process(self, data: Any, mode: str):
        self.index += 1

        tokenizer: PreTrainedTokenizer
        config = self.config
        max_seq_length = self.max_seq_length_dict[mode]
        tokenizer = self.tokenizer

        pair_data = data

        if "sptoken" not in data_conf:
            data_conf["sptoken"] = tokenizer.encode("",add_special_tokens=True)

        d = TokenIdsMakerForGLM.process(pair_data, tokenizer, max_seq_length,**data_conf)
        if self.index < 3:
            print(d)
        return d

    def collate_fn(self, batch):
        o = {k: [] for k in batch[0].keys()}
        for i, b in enumerate(batch):
            for k in b:
                o[k].append(torch.tensor(b[k]))

        ctxlens = o.pop('ctxlen')
        ctxlens2 = o.pop('ctxlen2')
        assert ctxlens is not None and ctxlens2 is not None

        seqlen = np.max([len(_) for _ in o['input_ids']])
        if 'input_ids2' in o:
            seqlen = np.max([seqlen] + [len(_) for _ in o['input_ids2']])

        tokenizer: PreTrainedTokenizer = self.tokenizer
        for k, v in o.items():
            pad_val = tokenizer.pad_token_id if 'label' not in k else -100
            o[k] = torch.stack(
                [F.pad(_, (0, seqlen - len(_)), mode='constant', value=pad_val) for _ in v])
        max_len = seqlen

        input_ids = o['input_ids']
        attention_mask, position_ids = build_masks_and_position_ids_glm(input_ids,ctxlens, max_len)
        o['attention_mask'] = attention_mask.bool()
        o['position_ids'] = position_ids.long()
        o["labels"] = o["labels"].long()

        input_ids2 = o['input_ids2']
        attention_mask, position_ids = build_masks_and_position_ids_glm(input_ids2,ctxlens2, max_len)
        o['attention_mask2'] = attention_mask.bool()
        o['position_ids2'] = position_ids.long()
        o["labels2"] = o["labels2"].long()
        return o
data_factory/data_helper_chatglm.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from typing import Any
import numpy as np
import torch
from aigc_zoo.model_zoo.chatglm2.llm_model import build_masks_and_position_ids_glm
from torch.nn import functional as F
from transformers import PreTrainedTokenizer
from data_factory.data_helper_base import NN_DataHelper_Base, data_conf
from data_factory.data_processer import TokenIdsMakerForGLM2


class NN_DataHelper_chatglm2(NN_DataHelper_Base):
    def on_data_process(self, data: Any, mode: str):
        self.index += 1

        tokenizer: PreTrainedTokenizer
        config = self.config
        max_seq_length = self.max_seq_length_dict[mode]
        tokenizer = self.tokenizer

        pair_data = data

        if "sptoken" not in data_conf:
            data_conf["sptoken"] = tokenizer.encode("", add_special_tokens=True)

        d = TokenIdsMakerForGLM2.process(pair_data, tokenizer, max_seq_length, **data_conf)
        if self.index < 3:
            print(d)
        return d

    def collate_fn(self, batch):
        o = {k: [] for k in batch[0].keys()}
        for i, b in enumerate(batch):
            for k in b:
                o[k].append(torch.tensor(b[k]))


        seqlen = np.max([len(_) for _ in o['input_ids']])
        if 'input_ids2' in o:
            seqlen = np.max([seqlen] + [len(_) for _ in o['input_ids2']])

        tokenizer: PreTrainedTokenizer = self.tokenizer
        for k, v in o.items():
            pad_val = tokenizer.pad_token_id if 'label' not in k else -100
            o[k] = torch.stack(
                [F.pad(_, (0, seqlen - len(_)), mode='constant', value=pad_val) for _ in v])
        max_len = seqlen

        input_ids = o['input_ids']
        attention_mask, position_ids = build_masks_and_position_ids_glm(input_ids,  max_len)
        o['attention_mask'] = attention_mask.bool()
        o['position_ids'] = position_ids.long()
        o['labels'] = o["labels"].long()

        input_ids2 = o['input_ids2']
        attention_mask, position_ids = build_masks_and_position_ids_glm(input_ids2,  max_len)
        o['attention_mask2'] = attention_mask.bool()
        o['position_ids2'] = position_ids.long()
        o['labels2'] = o["labels2"].long()
        return o
data_factory/data_helper_chatglm2.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_gpt2(NN_DataHelper_Base):...
data_factory/data_helper_gpt2.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_internlm(NN_DataHelper_Base):...
data_factory/data_helper_internlm.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_llama(NN_DataHelper_Base):...
data_factory/data_helper_llama.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/21 13:11

from data_factory.data_helper_base import NN_DataHelper_Base
from data_factory.data_helper_baichuan import NN_DataHelper_baichuan
from data_factory.data_helper_chatglm import NN_DataHelper_chatglm
from data_factory.data_helper_chatglm2 import NN_DataHelper_chatglm2
from data_factory.data_helper_bloom import NN_DataHelper_bloom
from data_factory.data_helper_internlm import NN_DataHelper_internlm
from data_factory.data_helper_gpt2 import NN_DataHelper_gpt2
from data_factory.data_helper_llama import NN_DataHelper_llama
from data_factory.data_helper_moss import NN_DataHelper_moss
from data_factory.data_helper_moss_plugin import NN_DataHelper_moss_plugin
from data_factory.data_helper_openbuddy import NN_DataHelper_openbuddy
from data_factory.data_helper_opt import NN_DataHelper_opt
from data_factory.data_helper_qwen import NN_DataHelper_qwen
from data_factory.data_helper_rwkv import NN_DataHelper_rwkv
from data_factory.data_helper_tiger import NN_DataHelper_tiger
from data_factory.data_helper_xverse import NN_DataHelper_xverse
data_factory/data_helper_loader.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_moss(NN_DataHelper_Base):...
data_factory/data_helper_moss.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_moss_plugin(NN_DataHelper_Base):...
data_factory/data_helper_moss_plugin.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_openbuddy(NN_DataHelper_Base):...
data_factory/data_helper_openbuddy.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_opt(NN_DataHelper_Base):...
data_factory/data_helper_opt.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_qwen(NN_DataHelper_Base):...
data_factory/data_helper_qwen.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_rwkv(NN_DataHelper_Base):...
data_factory/data_helper_rwkv.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_tiger(NN_DataHelper_Base):...
data_factory/data_helper_tiger.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/22 9:06
from data_factory.data_helper_base import NN_DataHelper_Base


class NN_DataHelper_xverse(NN_DataHelper_Base):...
data_factory/data_helper_xverse.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Time    : 2023/4/20 8:32
import copy
import json
import typing
import numpy as np
from transformers import PreTrainedTokenizer

DEFAULT_PAD_TOKEN = "[PAD]"
DEFAULT_EOS_TOKEN = "</s>"
DEFAULT_BOS_TOKEN = "</s>"
DEFAULT_UNK_TOKEN = "</s>"



class CorpusPreprocess:
   
    @classmethod
    def process(cls,tokenizer,lines):
        D = []
        for i, line in enumerate(lines):
            jd = json.loads(line)
            if not jd:
                continue
            prompt = jd['prompt']
            # response = jd['response']
            chosen = jd['chosen']
            rejected = jd['rejected']
            if chosen == rejected:
                continue
            D.append((prompt,chosen, rejected))
        return D



class TokenIdsMaker:
    @classmethod
    def trunction_ids(cls,a_ids: typing.List,b_ids: typing.List,max_seq_length,sptoken):
        while len(a_ids) + len(b_ids) > max_seq_length - len(sptoken) - 1:
            if len(a_ids) > len(b_ids):
                a_ids.pop(0)
            else:
                b_ids.pop(-1)

    @classmethod
    def process(cls, pair_data, tokenizer: PreTrainedTokenizer, sptoken, max_seq_length: int, src_max_length,
                dst_max_length):
        prompt, chosen, rejected = pair_data

        a_ids = tokenizer.encode(prompt, truncation=True, max_length=max_seq_length, add_special_tokens=False)
        b_ids1 = tokenizer.encode(chosen, truncation=True, max_length=max_seq_length, add_special_tokens=False)
        b_ids2 = tokenizer.encode(rejected, truncation=True, max_length=max_seq_length, add_special_tokens=False)

        if src_max_length is not None and src_max_length > 0:
            a_ids = a_ids[:src_max_length]

        if dst_max_length is not None and dst_max_length > 0:
            b_ids1 = b_ids1[:dst_max_length]
            b_ids2 = b_ids2[:dst_max_length]

        a_ids1 = copy.deepcopy(a_ids)
        cls.trunction_ids(a_ids1, b_ids1, max_seq_length, sptoken)

        a_ids2 = copy.deepcopy(a_ids)
        cls.trunction_ids(a_ids2, b_ids2, max_seq_length, sptoken)

        input_ids_a = np.asarray(sptoken + a_ids1 + b_ids1 + [tokenizer.eos_token_id], dtype=np.int32)
        attention_mask_a = np.asarray([1] * len(input_ids_a), dtype=np.int32)

        input_ids_b = np.asarray(sptoken + a_ids2 + b_ids2 + [tokenizer.eos_token_id], dtype=np.int32)
        attention_mask_b = np.asarray([1] * len(input_ids_b), dtype=np.int32)

        seqlen_a = len(input_ids_a)
        seqlen_b = len(input_ids_b)

        if seqlen_a == seqlen_b:
            if np.all(input_ids_a == input_ids_b):
                return None

        pos_a = len(a_ids1) + len(sptoken)
        pos_b = len(a_ids2) + len(sptoken)
        assert pos_a >= 0 and pos_a < max_seq_length - 1 and pos_b >= 0 and pos_b < max_seq_length - 1
        labels = np.asarray([-100] * pos_a + input_ids_a[pos_a:].tolist(), dtype=np.int64)
        labels2 = np.asarray([-100] * pos_b + input_ids_b[pos_b:].tolist(), dtype=np.int64)
        return {
            "input_ids": input_ids_a,
            "attention_mask": attention_mask_a,
            "labels": labels,
            # "seqlen": np.asarray(seqlen_a,dtype=np.int32),
            "input_ids2": input_ids_b,
            "attention_mask2": attention_mask_b,
            "labels2": labels2,
            # "seqlen2": np.asarray(seqlen_b, dtype=np.int32),
        }





class TokenIdsMakerForGLM(TokenIdsMaker):
    @classmethod
    def process(cls,pair_data,tokenizer: PreTrainedTokenizer,sptoken, max_seq_length: int,src_max_length,dst_max_length):
        prompt, chosen, rejected = pair_data


        a_ids = tokenizer.encode(prompt, truncation=True, max_length=max_seq_length,add_special_tokens=False)
        b_ids1 = tokenizer.encode(chosen, truncation=True, max_length=max_seq_length,add_special_tokens=False)
        b_ids2 = tokenizer.encode(rejected, truncation=True, max_length=max_seq_length,add_special_tokens=False)

        if src_max_length is not None and src_max_length > 0:
            a_ids = a_ids[:src_max_length]

        if dst_max_length is not None and dst_max_length > 0:
            b_ids1 = b_ids1[:dst_max_length]
            b_ids2 = b_ids2[:dst_max_length]

        a_ids1 = copy.deepcopy(a_ids)
        cls.trunction_ids(a_ids1,b_ids1,max_seq_length,sptoken)

        a_ids2 = copy.deepcopy(a_ids)
        cls.trunction_ids(a_ids2, b_ids2,max_seq_length,sptoken)


        input_ids_a = np.asarray( a_ids1 + sptoken + b_ids1 + [tokenizer.eos_token_id] ,dtype=np.int32)
        input_ids_b = np.asarray( a_ids2 + sptoken + b_ids2 + [tokenizer.eos_token_id] ,dtype=np.int32)


        seqlen_a = len(input_ids_a)
        seqlen_b = len(input_ids_b)

        if seqlen_a == seqlen_b:
            if np.all(input_ids_a == input_ids_b):
                return None

        pos_a = len(a_ids1)
        pos_b = len(a_ids2)
        assert pos_a >= 0 and pos_a < max_seq_length -1 and pos_b >= 0 and pos_b < max_seq_length -1
        labels = np.asarray([-100] * pos_a + input_ids_a[pos_a:].tolist(),dtype=np.int64)
        labels2 = np.asarray([-100] * pos_b + input_ids_b[pos_b:].tolist(),dtype=np.int64)

        ctxlen = pos_a + len(sptoken) - 1
        ctxlen2 = pos_b + len(sptoken) - 1

        ctxlen = np.asarray(ctxlen,dtype=np.int32)
        ctxlen2 = np.asarray(ctxlen2, dtype=np.int32)
        return {
            "input_ids": input_ids_a,
            "ctxlen": ctxlen,
            "labels": labels,
            # "seqlen": np.asarray(seqlen_a,dtype=np.int32),
            "input_ids2": input_ids_b,
            "ctxlen2": ctxlen2,
            "labels2": labels2,
            # "seqlen2": np.asarray(seqlen_b, dtype=np.int32),
        }


class TokenIdsMakerForGLM2(TokenIdsMaker):

    @classmethod
    def process(cls,pair_data,tokenizer: PreTrainedTokenizer,sptoken, max_seq_length: int,src_max_length,dst_max_length):
        prompt, chosen, rejected = pair_data


        a_ids = tokenizer.encode(prompt, truncation=True, max_length=max_seq_length,add_special_tokens=False)
        b_ids1 = tokenizer.encode(chosen, truncation=True, max_length=max_seq_length,add_special_tokens=False)
        b_ids2 = tokenizer.encode(rejected, truncation=True, max_length=max_seq_length,add_special_tokens=False)

        if src_max_length is not None and src_max_length > 0:
            a_ids = a_ids[:src_max_length]

        if dst_max_length is not None and dst_max_length > 0:
            b_ids1 = b_ids1[:dst_max_length]
            b_ids2 = b_ids2[:dst_max_length]

        a_ids1 = copy.deepcopy(a_ids)
        cls.trunction_ids(a_ids1,b_ids1,max_seq_length,sptoken)

        a_ids2 = copy.deepcopy(a_ids)
        cls.trunction_ids(a_ids2, b_ids2,max_seq_length,sptoken)


        input_ids_a = np.asarray( sptoken + a_ids1 + b_ids1 + [tokenizer.eos_token_id] ,dtype=np.int32)
        input_ids_b = np.asarray( sptoken + a_ids2 + b_ids2 + [tokenizer.eos_token_id] ,dtype=np.int32)


        seqlen_a = len(input_ids_a) + len(sptoken)
        seqlen_b = len(input_ids_b) + len(sptoken)

        if seqlen_a == seqlen_b:
            if np.all(input_ids_a == input_ids_b):
                return None

        pos_a = len(a_ids1)
        pos_b = len(a_ids2)
        assert pos_a >= 0 and pos_a < max_seq_length -1 and pos_b >= 0 and pos_b < max_seq_length -1
        labels = np.asarray([-100] * pos_a + input_ids_a[pos_a:].tolist(),dtype=np.int64)
        labels2 = np.asarray([-100] * pos_b + input_ids_b[pos_b:].tolist(),dtype=np.int64)


        return {
            "input_ids": input_ids_a,
            "labels": labels,
            # "seqlen": np.asarray(seqlen_a,dtype=np.int32),
            "input_ids2": input_ids_b,
            "labels2": labels2,
            # "seqlen2": np.asarray(seqlen_b, dtype=np.int32),
        }
data_factory/data_processer.py
ssbuild/dpo_finetuning
# @Time    : 2023/4/19 23:02
# @Author  : tk
# @FileName: data_utils
import sys
import os
sys.path.append(os.path.abspath(os.path.dirname(__file__)))

from collections import OrderedDict

from deep_training.data_helper import ModelArguments, TrainingArguments, DataArguments, TrainingArgumentsHF, \
    TrainingArgumentsCL, TrainingArgumentsAC
from deep_training.nlp.models.petl import PetlArguments
from deep_training.nlp.models.petl.prompt import PromptArguments
from transformers import HfArgumentParser
from data_factory.data_helper_loader import (NN_DataHelper_Base,
                                             NN_DataHelper_baichuan,
                                             NN_DataHelper_chatglm,
                                             NN_DataHelper_chatglm2,
                                             NN_DataHelper_bloom,
                                             NN_DataHelper_internlm,
                                             NN_DataHelper_gpt2,
                                             NN_DataHelper_llama,
                                             NN_DataHelper_moss,
                                             NN_DataHelper_moss_plugin,
                                             NN_DataHelper_opt,
                                             NN_DataHelper_qwen,
                                             NN_DataHelper_tiger,
                                             NN_DataHelper_xverse,
                                             NN_DataHelper_rwkv,
                                             NN_DataHelper_openbuddy)

from config import *
from module_setup import global_model_card


def _find_data_helper():
    data_helper_mapper = OrderedDict({
        "baichuan": NN_DataHelper_baichuan,
        "chatglm2": NN_DataHelper_chatglm2,
        "chatglm": NN_DataHelper_chatglm,
        "xverse": NN_DataHelper_xverse,
        "qwen": NN_DataHelper_qwen,
        "gpt2": NN_DataHelper_gpt2,
        "llama": NN_DataHelper_llama,
        "internlm": NN_DataHelper_internlm,
        "opt": NN_DataHelper_opt,
        "bloom": NN_DataHelper_bloom,
        "tiger": NN_DataHelper_tiger,
        "moss": NN_DataHelper_moss,
        "moss_plugin": NN_DataHelper_moss_plugin,
        "rwkv": NN_DataHelper_rwkv,
        "openbuddy": NN_DataHelper_openbuddy,
        "default": NN_DataHelper_Base,
    })
    NN_DataHelper = None
    for k in data_helper_mapper:
        if k in global_model_card:
            if k == "moss" and "plugin" in global_model_card:
                k = "moss_plugin"
            NN_DataHelper = data_helper_mapper[ k ]
            break
    return NN_DataHelper


NN_DataHelper = _find_data_helper()
if NN_DataHelper is None:
    NN_DataHelper = NN_DataHelper_Base
    raise ValueError(f"{global_model_card} for data_helper is not implemented ")

if __name__ == '__main__':
    if global_args[ "trainer_backend" ] == "hf":
        parser = HfArgumentParser((ModelArguments, TrainingArgumentsHF, DataArguments, PetlArguments),
                                  conflict_handler='resolve')
        model_args, training_args, data_args, lora_args = parser.parse_dict(train_info_args,
                                                                                         allow_extra_keys=True, )
    elif global_args[ "trainer_backend" ] == "pl":
        parser = HfArgumentParser((ModelArguments, TrainingArguments, DataArguments, PetlArguments))
        model_args, training_args, data_args, _ = parser.parse_dict(train_info_args)
    elif global_args["trainer_backend"] == "ac":
        parser = HfArgumentParser((ModelArguments, TrainingArgumentsCL, DataArguments, PetlArguments),
                                  conflict_handler='resolve')
        model_args, training_args, data_args, lora_args = parser.parse_dict(train_info_args, allow_extra_keys=True, )
    else:
        parser = HfArgumentParser((ModelArguments, TrainingArgumentsAC, DataArguments, PetlArguments),
                                  conflict_handler='resolve')
        model_args, training_args, data_args, lora_args = parser.parse_dict(train_info_args, allow_extra_keys=True, )

    dataHelper = NN_DataHelper(model_args, training_args, data_args)
    tokenizer, config, _, _ = dataHelper.load_tokenizer_and_config(config_kwargs={"torch_dtype": torch.float16})

    # 缓存数据集
    # 检测是否存在 output/dataset_0-train.record ，不存在则制作数据集
    dataHelper.make_dataset_all()


    # def shuffle_records(record_filenames, outfile, compression_type='GZIP'):
    #     print('shuffle_records record...')
    #     options = RECORD.TFRecordOptions(compression_type=compression_type)
    #     dataset_reader = Loader.RandomDataset(record_filenames, options=options, with_share_memory=True)
    #     data_size = len(dataset_reader)
    #     all_example = []
    #     for i in tqdm(range(data_size), desc='load records'):
    #         serialized = dataset_reader[i]
    #         all_example.append(serialized)
    #     dataset_reader.close()
    #
    #     shuffle_idx = list(range(data_size))
    #     random.shuffle(shuffle_idx)
    #     writer = WriterObject(outfile, options=options)
    #     for i in tqdm(shuffle_idx, desc='shuffle record'):
    #         example = all_example[i]
    #         writer.write(example)
    #     writer.close()
    #
    #
    # # 对每个record 再次打乱
    # for filename in dataHelper.train_files:
    #     shuffle_records(filename, filename)

data_utils.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Time:  21:08
# @Author: tk
# @File：__init__.py

infer/__init__.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Time    : 2023/5/17 11:36
import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),'..')))

from aigc_zoo.utils.llm_generate import Generate
from config import get_deepspeed_config
import torch
from deep_training.data_helper import ModelArguments
from transformers import HfArgumentParser,AutoConfig,PreTrainedTokenizer
from data_utils import train_info_args, NN_DataHelper
from aigc_zoo.model_zoo.auto.llm_model import MyTransformer

deep_config = get_deepspeed_config()

if __name__ == '__main__':
    train_info_args['seed'] = None
    parser = HfArgumentParser((ModelArguments, ))
    (model_args, ) = parser.parse_dict(train_info_args,allow_extra_keys=True)

    tokenizer : PreTrainedTokenizer
    dataHelper = NN_DataHelper(model_args)
    tokenizer, _, _, _ = dataHelper.load_tokenizer_and_config()

    ckpt_dir = './best_ckpt'
    config = AutoConfig.from_pretrained(ckpt_dir)
    pl_model = MyTransformer(config=config, model_args=model_args)
    if deep_config is None:
        train_weight = './best_ckpt/best.pt'
    else:
        # 建议直接使用转换脚本命令 支持 deepspeed stage 0,1,2,3， 生成 ./best_ckpt/last.ckpt/best.pt 权重文件
        # cd best_ckpt/last.ckpt
        # python zero_to_fp32.py . best.pt
        train_weight = './best_ckpt/last.ckpt/best.pt'

    pl_model.load_sft_weight(train_weight)


    pl_model.eval().half().cuda()

    pl_model.requires_grad_(False)
    model = pl_model.get_llm_model()

    text_list = ["写一个诗歌，关于冬天",
                 "晚上睡不着应该怎么办",
                 "从南京到上海的路线",
                 ]
    for input in text_list:
        response = Generate.generate(model, query=input, tokenizer=tokenizer, max_length=512,
                                     eos_token_id=config.eos_token_id,
                                     do_sample=False, top_p=0.7, temperature=0.95, )
        print('input', input)
        print('output', response)
infer/infer_finetuning.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Time    : 2023/5/17 11:36
import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),'..')))

from aigc_zoo.utils.llm_generate import Generate
import torch
from deep_training.data_helper import ModelArguments, DataArguments
from transformers import HfArgumentParser,AutoConfig,PreTrainedTokenizer
from data_utils import train_info_args, NN_DataHelper
from aigc_zoo.model_zoo.auto.llm_model import MyTransformer,PetlArguments

if __name__ == '__main__':
    train_info_args['seed'] = None
    parser = HfArgumentParser((ModelArguments, ))
    model_args, data_args = parser.parse_dict(train_info_args,allow_extra_keys=True)

    tokenizer : PreTrainedTokenizer
    dataHelper = NN_DataHelper(model_args)
    tokenizer, _, _, _ = dataHelper.load_tokenizer_and_config()

    ckpt_dir = '../scripts/best_ckpt/last'
    config = AutoConfig.from_pretrained(ckpt_dir)
    lora_args = PetlArguments.from_pretrained(ckpt_dir)

    assert lora_args.inference_mode == True

    new_num_tokens = config.vocab_size
    if config.task_specific_params is not None and config.task_specific_params.get('vocab_size', None) is not None:
        config.vocab_size = config.task_specific_params['vocab_size']

    pl_model = MyTransformer(config=config, model_args=model_args, lora_args=lora_args,
                                   torch_dtype=config.torch_dtype,
                                   new_num_tokens=new_num_tokens,
                                   
                                   # # device_map="auto",
                                   # device_map={"": 0},
                                   )
    # 加载sft权重
    pl_model.load_sft_weight(ckpt_dir)

    pl_model.eval().half().cuda()

    enable_merge_weight = False

    if enable_merge_weight:
        # 合并lora 权重 保存
        pl_model.save_sft_weight(os.path.join(ckpt_dir, 'pytorch_model_merge.bin'),merge_lora_weight=True)
    else:

        pl_model.requires_grad_(False)
        model = pl_model.get_llm_model()

        text_list = ["写一个诗歌，关于冬天",
                     "晚上睡不着应该怎么办",
                     "从南京到上海的路线",
                     ]
        for input in text_list:
            response = Generate.generate(model, query=input, tokenizer=tokenizer, max_length=512,
                                         eos_token_id=config.eos_token_id,
                                         do_sample=False, top_p=0.7, temperature=0.95, )
            print('input', input)
            print('output', response)
infer/infer_lora_finetuning.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/19 16:22

import os
from deep_training.utils.hf import register_transformer_model, register_transformer_config, \
    register_transformer_tokenizer
from transformers import AutoModelForCausalLM
from deep_training.nlp.models.rellama.modeling_llama import LlamaForCausalLM
from config import train_info_args

__all__ = [
    "module_setup",
    "global_model_card"
]

global_model_card = (train_info_args["model_name_or_path"] or train_info_args["config_name"])
global_model_card = global_model_card.split('/')[-1].lower() if os.path.isdir(global_model_card) else \
    str(os.path.dirname(global_model_card)).split('/')[-1].lower()


def module_setup():
    register_transformer_model(LlamaForCausalLM, AutoModelForCausalLM)

    if "baichuan" in global_model_card:
        assert "7b" in global_model_card or "13b" in global_model_card
        if "baichuan2" in global_model_card:
            if "7b" in global_model_card:
                from aigc_zoo.model_zoo.baichuan.baichuan2_7b.llm_model import (MyBaichuanForCausalLM as LM_MODEL, PetlArguments,  # noqa
                                                                                LoraConfig, PetlModel,
                                                                                PromptArguments,
                                                                                BaichuanConfig,
                                                                                BaichuanTokenizer)
            else:
                from aigc_zoo.model_zoo.baichuan.baichuan2_13b.llm_model import (MyBaichuanForCausalLM as LM_MODEL, PetlArguments,  # noqa
                                                                                 LoraConfig, PetlModel,
                                                                                 PromptArguments,
                                                                                 BaichuanConfig,
                                                                                 BaichuanTokenizer)
        else:
            if "7b" in global_model_card:
                from aigc_zoo.model_zoo.baichuan.baichuan_7b.llm_model import (MyBaichuanForCausalLM as LM_MODEL, PetlArguments,  # noqa
                                                                               LoraConfig, PetlModel,
                                                                               PromptArguments,
                                                                               BaichuanConfig,
                                                                               BaichuanTokenizer)
            else:
                from aigc_zoo.model_zoo.baichuan.baichuan_13b.llm_model import (MyBaichuanForCausalLM as LM_MODEL, PetlArguments,  # noqa
                                                                                LoraConfig, PetlModel,
                                                                                PromptArguments,
                                                                                BaichuanConfig,
                                                                                BaichuanTokenizer)

        register_transformer_config(BaichuanConfig)
        register_transformer_model(LM_MODEL, AutoModelForCausalLM)
        register_transformer_tokenizer(BaichuanConfig,BaichuanTokenizer,BaichuanTokenizer)

    elif "xverse" in global_model_card:
        from aigc_zoo.model_zoo.xverse.llm_model import (MyXverseForCausalLM as LM_MODEL,
                                                        PetlArguments,  # noqa
                                                        LoraConfig, PetlModel,
                                                        PromptArguments,
                                                        XverseConfig,)

        register_transformer_config(XverseConfig)
        register_transformer_model(LM_MODEL, AutoModelForCausalLM)

    elif "qwen" in global_model_card:
        from aigc_zoo.model_zoo.qwen.llm_model import (MyQWenLMHeadModel as LM_MODEL,
                                                     PetlArguments,  # noqa
                                                     LoraConfig, PetlModel,
                                                     PromptArguments,
                                                     QWenConfig,QWenTokenizer )

        register_transformer_config(QWenConfig)
        register_transformer_model(LM_MODEL, AutoModelForCausalLM)
        register_transformer_tokenizer(QWenConfig, QWenTokenizer, QWenTokenizer)

    elif "internlm" in global_model_card:
        from aigc_zoo.model_zoo.internlm.llm_model import (MyInternLMForCausalLM as LM_MODEL,
                                                     PetlArguments,  # noqa
                                                     LoraConfig, PetlModel,
                                                     PromptArguments,
                                                     InternLMConfig,InternLMTokenizer )

        register_transformer_config(InternLMConfig)
        register_transformer_model(LM_MODEL, AutoModelForCausalLM)
        register_transformer_tokenizer(InternLMConfig, InternLMTokenizer, InternLMTokenizer)

    elif "chatglm2" in global_model_card:
        from aigc_zoo.model_zoo.chatglm2.llm_model import (MyChatGLMForConditionalGeneration as LM_MODEL,
                                                           PetlArguments,  # noqa
                                                           LoraConfig, PetlModel,
                                                           PromptArguments,
                                                           ChatGLMConfig, ChatGLMTokenizer)

        register_transformer_config(ChatGLMConfig)
        register_transformer_model(LM_MODEL, AutoModelForCausalLM)
        register_transformer_tokenizer(ChatGLMConfig, ChatGLMTokenizer, ChatGLMTokenizer)

    elif "chatglm" in global_model_card:
        from aigc_zoo.model_zoo.chatglm.llm_model import (MyChatGLMForConditionalGeneration as LM_MODEL,
                                                           PetlArguments,  # noqa
                                                           LoraConfig, PetlModel,
                                                           PromptArguments,
                                                           ChatGLMConfig, ChatGLMTokenizer)

        register_transformer_config(ChatGLMConfig)
        register_transformer_model(LM_MODEL, AutoModelForCausalLM)
        register_transformer_tokenizer(ChatGLMConfig, ChatGLMTokenizer, ChatGLMTokenizer)

    elif "moss" in global_model_card:
        from aigc_zoo.model_zoo.moss.llm_model import (MyMossForCausalLM as LM_MODEL,
                                                      PetlArguments,  # noqa
                                                      LoraConfig, PetlModel,
                                                      PromptArguments,
                                                      MossConfig, MossTokenizer)

        register_transformer_config(MossConfig)
        register_transformer_model(LM_MODEL, AutoModelForCausalLM)
        register_transformer_tokenizer(MossConfig, MossTokenizer, MossTokenizer)
    elif "rwkv" in global_model_card:
        from aigc_zoo.model_zoo.rwkv4.llm_model import (MyRwkvForCausalLM as LM_MODEL,
                                                       PetlArguments,  # noqa
                                                       LoraConfig, PetlModel,
                                                       PromptArguments,
                                                       RwkvConfig)
        register_transformer_config(RwkvConfig)
        register_transformer_model(LM_MODEL, AutoModelForCausalLM)
        if 'world' in global_model_card:
            from aigc_zoo.model_zoo.rwkv4.rwkv4_tokenizer import RWKVTokenizer
            register_transformer_tokenizer(RwkvConfig, RWKVTokenizer, RWKVTokenizer)

    # 按需加入其他自定义模型
module_setup.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/10/12 10:50

import os
from config import global_args

def main():
    trainer_backend = global_args["trainer_backend"]
    if trainer_backend == "pl":
        from training.train_pl import main as main_execute
    elif trainer_backend == "hf":
        from training.train_hf import main as main_execute
    elif trainer_backend == "cl":
        from training.train_cl import main as main_execute
    elif trainer_backend == "ac":
        from training.train_ac import main as main_execute
    else:
        raise ValueError(f"{trainer_backend} NotImplemented ")

    main_execute()

def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


if __name__ == "__main__":
    main()

train.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Time:  21:07
# @Author: tk
# @File：__init__.py

training/__init__.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/25 12:29
import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),'..')))

import copy
import logging
import math
import datasets
import torch
import transformers
from deep_training.trainer.ac.trainer import TrainerAC
from transformers import (
    HfArgumentParser,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
from data_utils import NN_DataHelper, train_info_args, get_deepspeed_config, global_args
from aigc_zoo.model_zoo.auto.dpo_model import MyTransformerDPO,PetlArguments, LoraConfig
from deep_training.data_helper import ModelArguments, DataArguments,TrainingArgumentsAC

from module_setup import global_model_card

assert global_args["trainer_backend"] == "ac"

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.33.2")

logger = logging.getLogger(__name__)

# Setup logging
logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    handlers=[logging.StreamHandler(sys.stdout)],
)

def main():
    training_args: TrainingArgumentsAC
    parser = HfArgumentParser((ModelArguments, TrainingArgumentsAC, DataArguments, PetlArguments),
                              conflict_handler='resolve')
    model_args, training_args, data_args, lora_args = parser.parse_dict(train_info_args,allow_extra_keys=True,)
    lora_args = lora_args.config

    if training_args.should_log:
        # The default of training_args.log_level is passive, so we set log level at info here to have that default.
        transformers.utils.logging.set_verbosity_info()

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    dataHelper = NN_DataHelper(model_args, training_args, data_args)
    config_kwargs = {"torch_dtype": torch.float16}
    if global_args['config_merge']:
        config_kwargs.update(global_args['config_merge'])

    tokenizer, config, _, _ = dataHelper.load_tokenizer_and_config(config_kwargs=config_kwargs)

    with training_args.main_process_first(desc="make_dataset_all"):
        dataHelper.make_dataset_all()

    is_bf16_supported = torch.cuda.is_bf16_supported()
    precision = global_args[ "precision" ]
    if precision == "auto":
        # 精度 根据实际情况做调整
        if is_bf16_supported:
            precision = 'bf16'
        else:
            precision = '16'

        if global_args["quantization_config"] is not None and global_args["quantization_config"].load_in_8bit:
            precision = "32"


    if str(precision) == '16':
        training_args.fp16 = True
    elif str(precision) == 'bf16':
        training_args.bf16 = True
    else:
        training_args.fp16 = False
        training_args.bf16 = False

    if 'rwkv' in global_model_card:
        if precision.startswith('16'):
            RWKV_FLOAT_MODE = '16'
        elif precision.startswith('bf16'):
            RWKV_FLOAT_MODE = 'bf16'
        else:
            RWKV_FLOAT_MODE = '32'
        from deep_training.nlp.models.rwkv4.modeling_rwkv import set_model_profile
        # 加载cuda_core
        set_model_profile(RWKV_T_MAX=config.ctx_len, RWKV_FLOAT_MODE=RWKV_FLOAT_MODE)

    deepspeed_config = get_deepspeed_config(precision)
    if deepspeed_config:
        training_args.deepspeed = deepspeed_config

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    world_size,local_rank,process_index = training_args.world_size,training_args.local_rank,training_args.process_index

    dpo_args = dict(beta=0.1, ref_free=False)
    transformer_args = dict(config=config, model_args=model_args, training_args=training_args, lora_args=lora_args,
                            quantization_config=global_args["quantization_config"],
                            device_map={"": local_rank} if world_size > 1 else "auto",
                            torch_dtype=torch.float16,
                            new_num_tokens=len(tokenizer),  # 可能扩充词
                            **dpo_args
                            )

    if transformer_args["quantization_config"] is None:
        transformer_args.pop("device_map")

    pl_model = MyTransformerDPO(**transformer_args)

    pl_ref_model = copy.deepcopy(pl_model)
    pl_ref_model = pl_ref_model.eval().half()
    pl_ref_model.requires_grad_(False)

    pl_model.backbone.set_ref_model(pl_ref_model)

    config.save_pretrained(training_args.output_dir)

    # 加载sft权重
    # pl_model.load_sft_weight('./best_ckpt/best.pt',is_trainable=True)

    pl_model = pl_model.float()

    train_datasets = None
    if training_args.do_train:
        train_datasets = dataHelper.load_distributed_random_sampler(
            dataHelper.train_files,
            with_load_memory=data_args.data_backend == 'record',
            collate_fn=dataHelper.collate_fn,
            batch_size=training_args.per_device_train_batch_size,
            drop_last=training_args.dataloader_drop_last,  # 多卡建议扔掉
            num_processes=world_size, process_index=process_index,
            num_workers = training_args.dataloader_num_workers,
            pin_memory = training_args.dataloader_pin_memory,
        )



    # Initialize our Trainer
    trainer = TrainerAC(
        model=pl_model,
        args=training_args,
        train_dataset=train_datasets,
        tokenizer=tokenizer,
        # Data collator will default to DataCollatorWithPadding, so we change it.
        data_collator=default_data_collator,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        trainer.train(resume_from_checkpoint=checkpoint)




def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


if __name__ == "__main__":
    main()

training/train_ac.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/25 12:29
import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),'..')))

import copy
import logging
import math
from contextlib import nullcontext
import datasets
import torch
import transformers
from deep_training.trainer.cl.trainer import TrainerCL
from transformers import (
    HfArgumentParser,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
from data_utils import NN_DataHelper, train_info_args, get_deepspeed_config, global_args
from aigc_zoo.model_zoo.auto.dpo_model import MyTransformerDPO,PetlArguments, LoraConfig
from deep_training.data_helper import ModelArguments, DataArguments,TrainingArgumentsCL

from module_setup import global_model_card

assert global_args["trainer_backend"] == "cl"

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.33.2")


logger = logging.getLogger(__name__)

# Setup logging
logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    handlers=[logging.StreamHandler(sys.stdout)],
)

def main():
    world_size, local_rank, process_index = int(os.environ.get("WORLD_SIZE", 1)), int(
        os.environ.get("LOCAL_RANK", 0)), int(os.environ.get("RANK", 0))


    training_args: TrainingArgumentsCL
    parser = HfArgumentParser((ModelArguments, TrainingArgumentsCL, DataArguments, PetlArguments),
                              conflict_handler='resolve')
    model_args, training_args, data_args, lora_args = parser.parse_dict(train_info_args,allow_extra_keys=True,)
    lora_args = lora_args.config



    dataHelper = NN_DataHelper(model_args, training_args, data_args)
    config_kwargs = {"torch_dtype": torch.float16}
    if global_args['config_merge']:
        config_kwargs.update(global_args['config_merge'])

    tokenizer, config, _, _ = dataHelper.load_tokenizer_and_config(config_kwargs=config_kwargs)

    if process_index == 0:
        dataHelper.make_dataset_all()

    is_bf16_supported = torch.cuda.is_bf16_supported()
    precision = global_args[ "precision" ]
    if precision == "auto":
        # 精度 根据实际情况做调整
        if is_bf16_supported:
            precision = 'bf16'
        else:
            precision = '16'

        if global_args["quantization_config"] is not None and global_args["quantization_config"].load_in_8bit:
            precision = "32"


    if str(precision) == '16':
        training_args.fp16 = True
    elif str(precision) == 'bf16':
        training_args.bf16 = True
    else:
        training_args.fp16 = False
        training_args.bf16 = False

    if 'rwkv' in global_model_card:
        if precision.startswith('16'):
            RWKV_FLOAT_MODE = '16'
        elif precision.startswith('bf16'):
            RWKV_FLOAT_MODE = 'bf16'
        else:
            RWKV_FLOAT_MODE = '32'
        from deep_training.nlp.models.rwkv4.modeling_rwkv import set_model_profile
        # 加载cuda_core
        set_model_profile(RWKV_T_MAX=config.ctx_len, RWKV_FLOAT_MODE=RWKV_FLOAT_MODE)

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}"
        + f"16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    dpo_args = dict(beta=0.1, ref_free=False)
    transformer_args = dict(config=config, model_args=model_args, training_args=training_args, lora_args=lora_args,
                            quantization_config=global_args["quantization_config"],
                            device_map={"": local_rank} if world_size > 1 else "auto",
                            torch_dtype=torch.float16,
                            new_num_tokens=len(tokenizer),  # 可能扩充词
                            **dpo_args
                            )

    if transformer_args["quantization_config"] is None:
        transformer_args.pop("device_map")

    with nullcontext():
        pl_model = MyTransformerDPO(**transformer_args)

    pl_ref_model = copy.deepcopy(pl_model)
    pl_ref_model = pl_ref_model.eval().half()
    pl_ref_model.requires_grad_(False)

    pl_model.backbone.set_ref_model(pl_ref_model)

    config.save_pretrained(training_args.output_dir)

    # 加载sft权重
    # pl_model.load_sft_weight('./best_ckpt/best.pt',is_trainable=True)

    pl_model = pl_model.float()

    train_datasets = None
    if training_args.do_train:
        train_datasets = dataHelper.load_distributed_random_sampler(
            dataHelper.train_files,
            with_load_memory=data_args.data_backend == 'record',
            collate_fn=dataHelper.collate_fn,
            batch_size=training_args.per_device_train_batch_size,
            drop_last=training_args.dataloader_drop_last,  # 多卡建议扔掉
            num_processes=world_size, process_index=process_index,
            num_workers = training_args.dataloader_num_workers,
            pin_memory = training_args.dataloader_pin_memory,
        )



    # Initialize our Trainer
    trainer = TrainerCL(
        model=pl_model,
        args=training_args,
        train_dataset=train_datasets,
        tokenizer=tokenizer,
        # Data collator will default to DataCollatorWithPadding, so we change it.
        data_collator=default_data_collator,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        trainer.train(resume_from_checkpoint=checkpoint)




def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


if __name__ == "__main__":
    main()

training/train_cl.py
ssbuild/dpo_finetuning
# -*- coding: utf-8 -*-
# @Author  : ssbuild
# @Time    : 2023/9/25 12:29
import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),'..')))

import copy
import logging
import math
import datasets
import torch
import transformers
from deep_training.trainer.hf.trainer import TrainerHF
from transformers import (
    HfArgumentParser,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
from data_utils import NN_DataHelper, train_info_args, get_deepspeed_config, global_args
from aigc_zoo.model_zoo.auto.dpo_model import MyTransformerDPO,PetlArguments, LoraConfig
from deep_training.data_helper import ModelArguments, DataArguments,TrainingArgumentsHF

from module_setup import global_model_card

assert global_args["trainer_backend"] == "hf"

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.33.2")

logger = logging.getLogger(__name__)

# Setup logging
logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    handlers=[logging.StreamHandler(sys.stdout)],
)

def main():
    training_args: TrainingArgumentsHF
    parser = HfArgumentParser((ModelArguments, TrainingArgumentsHF, DataArguments, PetlArguments),
                              conflict_handler='resolve')
    model_args, training_args, data_args, lora_args = parser.parse_dict(train_info_args,allow_extra_keys=True,)
    lora_args = lora_args.config

    if training_args.should_log:
        # The default of training_args.log_level is passive, so we set log level at info here to have that default.
        transformers.utils.logging.set_verbosity_info()

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    dataHelper = NN_DataHelper(model_args, training_args, data_args)
    config_kwargs = {"torch_dtype": torch.float16}
    if global_args['config_merge']:
        config_kwargs.update(global_args['config_merge'])

    tokenizer, config, _, _ = dataHelper.load_tokenizer_and_config(config_kwargs=config_kwargs)

    with training_args.main_process_first(desc="make_dataset_all"):
        dataHelper.make_dataset_all()

    is_bf16_supported = torch.cuda.is_bf16_supported()
    precision = global_args[ "precision" ]
    if precision == "auto":
        # 精度 根据实际情况做调整
        if is_bf16_supported:
            precision = 'bf16'
        else:
            precision = '16'

        if global_args["quantization_config"] is not None and global_args["quantization_config"].load_in_8bit:
            precision = "32"


    if str(precision) == '16':
        training_args.fp16 = True
    elif str(precision) == 'bf16':
        training_args.bf16 = True
    else:
        training_args.fp16 = False
        training_args.bf16 = False

    if 'rwkv' in global_model_card:
        if precision.startswith('16'):
            RWKV_FLOAT_MODE = '16'
        elif precision.startswith('bf16'):
            RWKV_FLOAT_MODE = 'bf16'
        else:
            RWKV_FLOAT_MODE = '32'
        from deep_training.nlp.models.rwkv4.modeling_rwkv import set_model_profile
        # 加载cuda_core
        set_model_profile(RWKV_T_MAX=config.ctx_len, RWKV_FLOAT_MODE=RWKV_FLOAT_MODE)

    deepspeed_config = get_deepspeed_config(precision)
    if deepspeed_config:
        training_args.deepspeed = deepspeed_config

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    world_size,local_rank,process_index = training_args.world_size,training_args.local_rank,training_args.process_index

    dpo_args = dict(beta=0.1, ref_free=False)
    transformer_args = dict(config=config, model_args=model_args, training_args=training_args, lora_args=lora_args,
                            quantization_config=global_args["quantization_config"],
                            device_map={"": local_rank} if world_size > 1 else "auto",
                            torch_dtype=torch.float16,
                            new_num_tokens=len(tokenizer),  # 可能扩充词
                            **dpo_args
                            )

    if transformer_args["quantization_config"] is None:
        transformer_args.pop("device_map")

    pl_model = MyTransformerDPO(**transformer_args)

    pl_ref_model = copy.deepcopy(pl_model)
    pl_ref_model = pl_ref_model.eval().half()
    pl_ref_model.requires_grad_(False)

    pl_model.backbone.set_ref_model(pl_ref_model)

    config.save_pretrained(training_args.output_dir)

    # 加载sft权重
    # pl_model.load_sft_weight('./best_ckpt/best.pt',is_trainable=True)

    pl_model = pl_model.float()

    train_datasets = None
    if training_args.do_train:
        train_datasets = dataHelper.load_distributed_random_sampler(
            dataHelper.train_files,
            with_load_memory=data_args.data_backend == 'record',
            collate_fn=dataHelper.collate_fn,
            batch_size=training_args.per_device_train_batch_size,
            drop_last=training_args.dataloader_drop_last,  # 多卡建议扔掉
            num_processes=world_size, process_index=process_index,
            num_workers = training_args.dataloader_num_workers,
            pin_memory = training_args.dataloader_pin_memory,
        )



    # Initialize our Trainer
    trainer = TrainerHF(
        model=pl_model,
        args=training_args,
        train_dataset=train_datasets,
        tokenizer=tokenizer,
        # Data collator will default to DataCollatorWithPadding, so we change it.
        data_collator=default_data_collator,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()  # Saves the tokenizer too for easy upload

        metrics = train_result.metrics
        metrics["train_samples"] = len(train_datasets)
        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()




def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


if __name__ == "__main__":
    main()

training/train_hf.py
ssbuild/dpo_finetuning
# @Time    : 2023/4/19 23:03
# @Author  : tk
# @FileName: train.py
import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),'..')))

import copy
import torch
from deep_training.data_helper import ModelArguments, DataArguments, TrainingArguments
from deep_training.trainer.pl.modelcheckpoint import ModelCheckpointEx
from lightning import Trainer
from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor
from lightning.pytorch.strategies import DeepSpeedStrategy
from transformers import HfArgumentParser
from data_utils import NN_DataHelper, train_info_args, global_args, get_deepspeed_config
from aigc_zoo.model_zoo.auto.dpo_model import MyTransformerDPO,PetlArguments, LoraConfig
from module_setup import global_model_card

assert global_args["trainer_backend"] == "pl"

def main():
    parser = HfArgumentParser((ModelArguments, TrainingArguments, DataArguments, PetlArguments))
    model_args, training_args, data_args, lora_args = parser.parse_dict(train_info_args)
    lora_args = lora_args.config

    output_weight_dir = './best_ckpt'

    dataHelper = NN_DataHelper(model_args, training_args, data_args)
    config_kwargs = {}
    if global_args['config_merge']:
        config_kwargs.update(global_args['config_merge'])
    tokenizer, config, _, _ = dataHelper.load_tokenizer_and_config(config_kwargs=config_kwargs)


    dataHelper.make_dataset_all()

    is_bf16_supported = torch.cuda.is_bf16_supported()
    precision = global_args[ "precision" ]
    if precision == "auto":
        # 精度 根据实际情况做调整
        if is_bf16_supported:
            precision = 'bf16'
        else:
            precision = '16'

        if global_args["quantization_config"] is not None and global_args["quantization_config"].load_in_8bit:
            precision = "32"

    if 'rwkv' in global_model_card:
        if precision.startswith('16'):
            RWKV_FLOAT_MODE = '16'
        elif precision.startswith('bf16'):
            RWKV_FLOAT_MODE = 'bf16'
        else:
            RWKV_FLOAT_MODE = '32'
        from deep_training.nlp.models.rwkv4.modeling_rwkv import set_model_profile
        # 加载cuda_core
        set_model_profile(RWKV_T_MAX=config.ctx_len, RWKV_FLOAT_MODE=RWKV_FLOAT_MODE)

    deepspeed_config = get_deepspeed_config(precision)
    strategy = 'ddp' if torch.cuda.device_count() > 1 else 'auto'
    if deepspeed_config is not None and len(deepspeed_config):
        strategy = DeepSpeedStrategy(config=deepspeed_config, )
    checkpoint_callback = ModelCheckpointEx(
        # monitor='loss',
        dirpath=output_weight_dir,
        save_weights_only=True,
        save_last=True,
        save_top_k=1,
        # every_n_train_steps=2000 // training_args.gradient_accumulation_steps,
        every_n_epochs=1,
        lora_args=lora_args,
    )


    trainer = Trainer(
        callbacks=[checkpoint_callback, LearningRateMonitor(logging_interval='step')],
        max_epochs=training_args.max_epochs,
        max_steps=training_args.max_steps,
        accelerator="gpu",
        devices=data_args.devices,
        enable_progress_bar=True,
        default_root_dir=data_args.output_dir,
        gradient_clip_val=training_args.max_grad_norm,
        accumulate_grad_batches=training_args.gradient_accumulation_steps,
        num_sanity_val_steps=0,
        strategy=strategy,
        precision=precision,# 可以自行尝试  "32": "32-true", "16": "16-mixed", "bf16": "bf16-mixed"
    )


    dpo_args = dict(beta=0.1, ref_free=False)
    transformer_args = dict(config=config, model_args=model_args, training_args=training_args, lora_args=lora_args,
                            quantization_config=global_args["quantization_config"],
                            device_map={"": trainer.local_rank} if trainer.world_size > 1 else "auto",
                            torch_dtype=torch.float16,
                            new_num_tokens=len(tokenizer),  # 可能扩充词
                            **dpo_args
                           )
    # 移除device_map
    if global_args["quantization_config"] is None:
        transformer_args.pop("device_map")


    pl_model = MyTransformerDPO(**transformer_args)

    pl_ref_model = copy.deepcopy(pl_model)
    pl_ref_model = pl_ref_model.eval().half()
    pl_ref_model.requires_grad_(False)

    pl_model.backbone.set_ref_model(pl_ref_model)

    config.save_pretrained(output_weight_dir)

    # 如果自定义训练了sft_weight , 可以再次加载sft_weight
    # pl_model.load_sft_weight('sft_weight.bin',is_trainable=True)

    pl_model = pl_model.float()


    train_datasets = dataHelper.load_distributed_random_sampler(
        dataHelper.train_files,
        with_load_memory=True,
        collate_fn=dataHelper.collate_fn,
        batch_size=training_args.train_batch_size,
        num_workers=0,  # num_workers for DataLoader
        drop_last=True,  # 多卡建议扔掉
        num_processes=trainer.world_size, process_index=trainer.global_rank)

    if train_datasets is not None:
        trainer.fit(pl_model, train_dataloaders=train_datasets)


def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


if __name__ == "__main__":
    main()
training/train_pl.py
ssbuild/dpo_finetuning