1import torch
2from datasets import load_dataset
3from transformers import AutoModelForCausalLM, AutoTokenizer
4import multiprocessing
5from trl import (
6 DPOConfig,
7 DPOTrainer,
8 ModelConfig,
9 ScriptArguments,
10 TrlParser,
11 get_kbit_device_map,
12 get_peft_config,
13 get_quantization_config,
14)
15from trl.trainer.utils import SIMPLE_CHAT_TEMPLATE
16
17if __name__ == "__main__":
18 parser = TrlParser((ScriptArguments, DPOConfig, ModelConfig))
19 script_args, training_args, model_config = parser.parse_args_and_config()
20
21 torch_dtype = (
22 model_config.torch_dtype
23 if model_config.torch_dtype in ["auto", None]
24 else getattr(torch, model_config.torch_dtype)
25 )
26
27 quantization_config = get_quantization_config(model_config)
28
29 model_kwargs = dict(
30 revision=model_config.model_revision,
31 attn_implementation=model_config.attn_implementation,
32 torch_dtype=torch_dtype,
33 use_cache=False if training_args.gradient_checkpointing else True,
34 device_map=get_kbit_device_map() if quantization_config is not None else None,
35 quantization_config=quantization_config,
36 )
37
38 model = AutoModelForCausalLM.from_pretrained(
39 model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs
40 )
41
42 peft_config = get_peft_config(model_config)
43 if peft_config is None:
44 ref_model = AutoModelForCausalLM.from_pretrained(
45 model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs
46 )
47 else:
48 ref_model = None
49
50 tokenizer = AutoTokenizer.from_pretrained(
51 model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code
52 )
53 if tokenizer.pad_token is None:
54 tokenizer.pad_token = tokenizer.eos_token
55 if tokenizer.chat_template is None:
56 tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE
57 if script_args.ignore_bias_buffers:
58 model._ddp_params_and_buffers_to_ignore = [
59 name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool
60 ]
61
62 dataset = load_dataset(script_args.dataset_name,
63 split=script_args.dataset_train_split)
64 dataset=dataset.select_columns(['chosen', 'prompt', 'rejected'])
65
66 trainer = DPOTrainer(
67 model,
68 ref_model,
69 args=training_args,
70 train_dataset=dataset,
71 processing_class=tokenizer,
72 peft_config=peft_config,
73 )
74
75 trainer.train()
76
77 trainer.save_model(training_args.output_dir)
