LLMs_train/config.py at main · 5663015/LLMs_train · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
from transformers import (
    AutoModel,
    AutoModelForCausalLM,
    AutoTokenizer,
    LlamaForCausalLM,
    LlamaTokenizer,
    BloomForCausalLM,
    BloomTokenizerFast,
    GPTNeoXForCausalLM,
    GPTNeoXTokenizerFast,
)


class CONFIG:
    # 模型
    MODEL_MAP = {
        'llama': LlamaForCausalLM,
        'glm': AutoModel,
        'bloom': BloomForCausalLM,
        'pythia': GPTNeoXForCausalLM,
        'baichuan': AutoModelForCausalLM
    }

    # 分词器
    TOKENIZER_MAP = {
        'llama': LlamaTokenizer,
        'glm': AutoTokenizer,
        'bloom': BloomTokenizerFast,
        'pythia': GPTNeoXTokenizerFast,
        'baichuan': AutoTokenizer
    }

    # bos_id & eos_id & pad_id
    SPECIAL_IDS = {
        'BelleGroup-BELLE-LLaMA-EXT-13B': {'bos_id': 1, 'eos_id': 2, 'pad_id': 0},
        'open_llama_7b': {'bos_id': 1, 'eos_id': 2, 'pad_id': 0},
        'open_llama_13b': {'bos_id': 1, 'eos_id': 2, 'pad_id': 0},
        'chinese-alpaca-13b': {'bos_id': 1, 'eos_id': 2, 'pad_id': 49953},
        'chatglm2-6b': {'bos_id': None, 'eos_id': 2, 'pad_id': 0},
        'baichuan-7B': {'bos_id': 1, 'eos_id': 2, 'pad_id': 0},
        'Baichuan-13B-Base': {'bos_id': 1, 'eos_id': 2, 'pad_id': 0},
        'Baichuan-13B-Chat': {'bos_id': 1, 'eos_id': 2, 'pad_id': 0},
        'bloomz-7b1': {'bos_id': 1, 'eos_id': 2, 'pad_id': 3},
        'bloomz-1b7': {'bos_id': 1, 'eos_id': 2, 'pad_id': 3},
        'tigerbot-7b-sft': {'bos_id': 1, 'eos_id': 2, 'pad_id': 3},
        'tigerbot-7b-base': {'bos_id': 1, 'eos_id': 2, 'pad_id': 3},
        'pythia-12b-deduped': {'bos_id': 0, 'eos_id': 0, 'pad_id': 1},
        'pythia-12b-sft-v8-7k-steps': {'bos_id': 0, 'eos_id': 0, 'pad_id': 1},
        'pythia-1b-deduped': {'bos_id': 0, 'eos_id': 0, 'pad_id': 1},
    }

    # LORA配置
    LORA_MAP = {
        'llama': {
            "lora_r": 8,
            "lora_alpha": 32,
            "lora_dropout": 0.05,
            "lora_target_modules": "q_proj,v_proj,k_proj,o_proj,gate_proj,down_proj,up_proj",
            "modules_to_save": "null"
        },
        'glm': {
            "lora_r": 8,
            "lora_alpha": 32,
            "lora_dropout": 0.05,
            "lora_target_modules": "query_key_value,dense,dense_h_to_4h,dense_4h_to_h",
            "modules_to_save": "null"
        },
        'bloom': {
            "lora_r": 8,
            "lora_alpha": 32,
            "lora_dropout": 0.05,
            "lora_target_modules": "query_key_value",
            "modules_to_save": "null"
        },
        'pythia': {
            "lora_r": 8,
            "lora_alpha": 32,
            "lora_dropout": 0.05,
            "lora_target_modules": "query_key_value",
            "modules_to_save": "null"
        },
        'baichuan': {
            "lora_r": 8,
            "lora_alpha": 32,
            "lora_dropout": 0.05,
            "lora_target_modules": "W_pack,o_proj",
            "modules_to_save": "null"
        }
    }