llm-adapters

1
import os
2
import sys
3
from typing import List
4

5
import fire
6
import torch
7
import transformers
8
from datasets import load_dataset
9
from typing import List, Optional, Union
10

11
"""
12
Unused imports:
13
import torch.nn as nn
14
import bitsandbytes as bnb
15
"""
16
sys.path.append(os.path.join(os.getcwd(), "peft/src/"))
17
from peft import (  # noqa: E402
18
    LoraConfig,
19
    BottleneckConfig,
20
    PrefixTuningConfig,
21
    get_peft_model,
22
    get_peft_model_state_dict,
23
    prepare_model_for_int8_training,
24
    set_peft_model_state_dict,
25
)
26
from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaTokenizer, AutoModel  # noqa: F402
27

28

29
def train(
30
        # model/data params
31
        base_model: str = "",  # the only required argument
32
        data_path: str = "yahma/alpaca-cleaned",
33
        output_dir: str = "./lora-alpaca",
34
        adapter_name: str = "lora",
35
        load_8bit : bool = False,
36
        # training hyperparams
37
        batch_size: int = 128,
38
        micro_batch_size: int = 4,
39
        num_epochs: int = 3,
40
        learning_rate: float = 3e-4,
41
        cutoff_len: int = 256,
42
        val_set_size: int = 2000,
43
        use_gradient_checkpointing: bool = False,
44
        eval_step: int = 200,
45
        save_step: int = 200,
46
        # lora hyperparams
47
        lora_r: int = 8,
48
        lora_alpha: int = 16,
49
        lora_dropout: float = 0.05,
50
        lora_target_modules: List[str] = None,
51
        # bottleneck adapter hyperparams
52
        bottleneck_size: int = 256,
53
        non_linearity: str = "tanh",
54
        adapter_dropout: float = 0.0,
55
        use_parallel_adapter: bool = False,
56
        use_adapterp: bool = False,
57
        target_modules: List[str] = None,
58
        scaling: Union[float, str] = 1.0,
59
        # prefix tuning hyperparams
60
        num_virtual_tokens: int = 30,
61
        # llm hyperparams
62
        train_on_inputs: bool = True,  # if False, masks out inputs in loss
63
        group_by_length: bool = False,  # faster, but produces an odd training loss curve
64
        # wandb params
65
        wandb_project: str = "",
66
        wandb_run_name: str = "",
67
        wandb_watch: str = "",  # options: false | gradients | all
68
        wandb_log_model: str = "",  # options: false | true
69
        resume_from_checkpoint: str = None,  # either training checkpoint or final adapter
70
):
71
    print(
72
        f"Finetuning model with params:\n"
73
        f"base_model: {base_model}\n"
74
        f"data_path: {data_path}\n"
75
        f"output_dir: {output_dir}\n"
76
        f"batch_size: {batch_size}\n"
77
        f"micro_batch_size: {micro_batch_size}\n"
78
        f"num_epochs: {num_epochs}\n"
79
        f"learning_rate: {learning_rate}\n"
80
        f"cutoff_len: {cutoff_len}\n"
81
        f"val_set_size: {val_set_size}\n"
82
        f"use_gradient_checkpointing: {use_gradient_checkpointing}\n"
83
        f"lora_r: {lora_r}\n"
84
        f"lora_alpha: {lora_alpha}\n"
85
        f"lora_dropout: {lora_dropout}\n"
86
        f"lora_target_modules: {lora_target_modules}\n"
87
        f"bottleneck_size: {bottleneck_size}\n"
88
        f"non_linearity: {non_linearity}\n"
89
        f"adapter_dropout: {adapter_dropout}\n"
90
        f"use_parallel_adapter: {use_parallel_adapter}\n"
91
        f"use_adapterp: {use_adapterp}\n"
92
        f"train_on_inputs: {train_on_inputs}\n"
93
        f"scaling: {scaling}\n"
94
        f"adapter_name: {adapter_name}\n"
95
        f"target_modules: {target_modules}\n"
96
        f"group_by_length: {group_by_length}\n"
97
        f"wandb_project: {wandb_project}\n"
98
        f"wandb_run_name: {wandb_run_name}\n"
99
        f"wandb_watch: {wandb_watch}\n"
100
        f"wandb_log_model: {wandb_log_model}\n"
101
        f"resume_from_checkpoint: {resume_from_checkpoint}\n"
102
    )
103
    assert (
104
        base_model
105
    ), "Please specify a --base_model, e.g. --base_model='decapoda-research/llama-7b-hf'"
106
    gradient_accumulation_steps = batch_size // micro_batch_size
107

108
    device_map = "auto"
109
    world_size = int(os.environ.get("WORLD_SIZE", 1))
110
    ddp = world_size != 1
111
    if ddp:
112
        device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)}
113
        gradient_accumulation_steps = gradient_accumulation_steps // world_size
114

115
    # Check if parameter passed or if set within environ
116
    use_wandb = len(wandb_project) > 0 or (
117
            "WANDB_PROJECT" in os.environ and len(os.environ["WANDB_PROJECT"]) > 0
118
    )
119
    # Only overwrite environ if wandb param passed
120
    if len(wandb_project) > 0:
121
        os.environ["WANDB_PROJECT"] = wandb_project
122
    if len(wandb_watch) > 0:
123
        os.environ["WANDB_WATCH"] = wandb_watch
124
    if len(wandb_log_model) > 0:
125
        os.environ["WANDB_LOG_MODEL"] = wandb_log_model
126

127
    if load_8bit:
128
        model = AutoModelForCausalLM.from_pretrained(
129
            base_model,
130
            load_in_8bit=load_8bit,
131
            torch_dtype=torch.float16,
132
            device_map=device_map,
133
            trust_remote_code=True,
134
        )
135
    else:
136
        model = AutoModelForCausalLM.from_pretrained(
137
            base_model,
138
            load_in_8bit=False,
139
            torch_dtype=torch.float16,
140
            device_map={"": int(os.environ.get("LOCAL_RANK") or 0)},
141
            trust_remote_code=True,
142
        )
143

144
    if model.config.model_type == "llama":
145
        # Due to the name of transformers' LlamaTokenizer, we have to do this
146
        tokenizer = LlamaTokenizer.from_pretrained(base_model)
147
    else:
148
        tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
149

150
    tokenizer.pad_token_id = (
151
        0  # unk. we want this to be different from the eos token
152
    )
153
    tokenizer.padding_side = "left"  # Allow batched inference
154

155
    def tokenize(prompt, add_eos_token=True):
156
        # there's probably a way to do this with the tokenizer settings
157
        # but again, gotta move fast
158
        result = tokenizer(
159
            prompt,
160
            truncation=True,
161
            max_length=cutoff_len,
162
            padding=False,
163
            return_tensors=None,
164
        )
165
        if (
166
                result["input_ids"][-1] != tokenizer.eos_token_id
167
                and len(result["input_ids"]) < cutoff_len
168
                and add_eos_token
169
        ):
170
            result["input_ids"].append(tokenizer.eos_token_id)
171
            if "chatglm" not in base_model:
172
                result["attention_mask"].append(1)
173

174
        result["labels"] = result["input_ids"].copy()
175

176
        if "chatglm" in base_model:
177
            return {"input_ids": result["input_ids"], "labels": result["labels"]}
178
        else:
179
            return result
180

181
    def generate_and_tokenize_prompt(data_point):
182
        full_prompt = generate_prompt(data_point)
183
        tokenized_full_prompt = tokenize(full_prompt)
184
        if not train_on_inputs:
185
            user_prompt = generate_prompt({**data_point, "output": ""})
186
            tokenized_user_prompt = tokenize(user_prompt, add_eos_token=False)
187
            user_prompt_len = len(tokenized_user_prompt["input_ids"])
188

189
            tokenized_full_prompt["labels"] = [
190
                                                  -100
191
                                              ] * user_prompt_len + tokenized_full_prompt["labels"][
192
                                                                    user_prompt_len:
193
                                                                    ]  # could be sped up, probably
194
        return tokenized_full_prompt
195

196
    model = prepare_model_for_int8_training(model, use_gradient_checkpointing=use_gradient_checkpointing)
197
    if adapter_name == "lora":
198
        config = LoraConfig(
199
            r=lora_r,
200
            lora_alpha=lora_alpha,
201
            target_modules=target_modules,
202
            lora_dropout=lora_dropout,
203
            bias="none",
204
            task_type="CAUSAL_LM",
205
        )
206
    elif adapter_name == "bottleneck":
207
        config = BottleneckConfig(
208
            bottleneck_size=bottleneck_size,
209
            non_linearity=non_linearity,
210
            adapter_dropout=adapter_dropout,
211
            use_parallel_adapter=use_parallel_adapter,
212
            use_adapterp=use_adapterp,
213
            target_modules=target_modules,
214
            scaling=scaling,
215
            bias="none",
216
            task_type="CAUSAL_LM",
217
        )
218
    elif adapter_name == "prefix-tuning":
219
        config = PrefixTuningConfig(
220
            num_virtual_tokens=num_virtual_tokens,
221
            task_type="CAUSAL_LM",
222
        )
223
    model = get_peft_model(model, config)
224
    if adapter_name == "prefix-tuning":
225
        model.to('cuda')
226

227
    if data_path.endswith(".json"):  # todo: support jsonl
228
        data = load_dataset("json", data_files=data_path)
229
    else:
230
        data = load_dataset(data_path)
231

232
    if resume_from_checkpoint:
233
        # Check the available weights and load them
234
        checkpoint_name = os.path.join(
235
            resume_from_checkpoint, "pytorch_model.bin"
236
        )  # Full checkpoint
237
        if not os.path.exists(checkpoint_name):
238
            checkpoint_name = os.path.join(
239
                resume_from_checkpoint, "adapter_model.bin"
240
            )  # only LoRA model - LoRA config above has to fit
241
            resume_from_checkpoint = (
242
                False  # So the trainer won't try loading its state
243
            )
244
        # The two files above have a different name depending on how they were saved, but are actually the same.
245
        if os.path.exists(checkpoint_name):
246
            print(f"Restarting from {checkpoint_name}")
247
            adapters_weights = torch.load(checkpoint_name)
248
            model = set_peft_model_state_dict(model, adapters_weights)
249
        else:
250
            print(f"Checkpoint {checkpoint_name} not found")
251

252
    model.print_trainable_parameters()  # Be more transparent about the % of trainable params.
253

254
    if val_set_size > 0:
255
        train_val = data["train"].train_test_split(
256
            test_size=val_set_size, shuffle=True, seed=42
257
        )
258
        train_data = (
259
            train_val["train"].shuffle().map(generate_and_tokenize_prompt)
260
        )
261
        val_data = (
262
            train_val["test"].shuffle().map(generate_and_tokenize_prompt)
263
        )
264
    else:
265
        train_data = data["train"].shuffle().map(generate_and_tokenize_prompt)
266
        val_data = None
267

268
    if not ddp and torch.cuda.device_count() > 1:
269
        # keeps Trainer from trying its own DataParallelism when more than 1 gpu is available
270
        model.is_parallelizable = True
271
        model.model_parallel = True
272

273
    trainer = transformers.Trainer(
274
        model=model,
275
        train_dataset=train_data,
276
        eval_dataset=val_data,
277
        args=transformers.TrainingArguments(
278
            per_device_train_batch_size=micro_batch_size,
279
            gradient_accumulation_steps=gradient_accumulation_steps,
280
            warmup_steps=100,
281
            num_train_epochs=num_epochs,
282
            learning_rate=learning_rate,
283
            fp16=True,
284
            logging_steps=10,
285
            optim="adamw_torch",
286
            evaluation_strategy="steps" if val_set_size > 0 else "no",
287
            save_strategy="steps",
288
            eval_steps=eval_step if val_set_size > 0 else None,
289
            save_steps=save_step,
290
            output_dir=output_dir,
291
            save_total_limit=3,
292
            load_best_model_at_end=True if val_set_size > 0 else False,
293
            ddp_find_unused_parameters=False if ddp else None,
294
            group_by_length=group_by_length,
295
            report_to="wandb" if use_wandb else None,
296
            run_name=wandb_run_name if use_wandb else None,
297
        ),
298
        data_collator=transformers.DataCollatorForSeq2Seq(
299
            tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True
300
        ),
301
    )
302
    model.config.use_cache = False
303

304
    old_state_dict = model.state_dict
305
    model.state_dict = (
306
        lambda self, *_, **__: get_peft_model_state_dict(
307
            self, old_state_dict()
308
        )
309
    ).__get__(model, type(model))
310

311
    if torch.__version__ >= "2" and sys.platform != "win32":
312
        model = torch.compile(model)
313

314
    trainer.train(resume_from_checkpoint=resume_from_checkpoint)
315

316
    model.save_pretrained(output_dir)
317

318
    print(
319
        "\n If there's a warning about missing keys above, please disregard :)"
320
    )
321

322

323
def generate_prompt(data_point):
324
    # sorry about the formatting disaster gotta move fast
325
    if data_point["input"]:
326
        return f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request. 
327

328
                ### Instruction:
329
                {data_point["instruction"]}
330
                
331
                ### Input:
332
                {data_point["input"]}
333
                
334
                ### Response:
335
                {data_point["output"]}""" # noqa: E501
336
    else:
337
        return f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.  
338

339
                ### Instruction:
340
                {data_point["instruction"]}
341
                
342
                ### Response:
343
                {data_point["output"]}""" # noqa: E501
344

345

346
if __name__ == "__main__":
347
    fire.Fire(train)
348