├── DPO_train.sh
├── README.md
├── merge.py
└── training_DPO.py


/DPO_train.sh:
--------------------------------------------------------------------------------
 1 | # 내가 훈련할 모델의 레포입니다.
 2 | BASIS=MODEL_REPO
 3 | 
 4 | # 내가 저장한 데이터의 repo입니다
 5 | DATA=Your_DPO_Dataset
 6 | 
 7 | # DPO가 진행된 후 저장되는 폴더의 명칭입니다.
 8 | output_directory=save_folder_name
 9 | 
10 | # 훈련이 끝난 어댑터와, 기존의 SFT를 merge한 이후 저장되는 폴더 이름입니다.
11 | final_dir=final_merged_model
12 | 
13 | 
14 | 
15 | # accelerate을 통해 분산학습을 시도하셔도 됩니다.
16 | python training_DPO.py \
17 |     --model_name_or_path $BASIS \
18 |     --output_dir $output_directory \
19 |     --datapath $DATA \
20 |     --num_epochs 1 \
21 |     --lora_r 32 \
22 |     --lora_alpha 16 \
23 |     --lora_dropout 0.01 \
24 |     --per_device_train_batch_size 64 \
25 |     --lr_scheduler_type 'cosine' \
26 |     --gradient_accumulation_steps 4 \
27 |     --eval_step 100 \
28 |     --max_prompt_length 8192 \
29 |     --max_length 8192 \
30 | 
31 | 
32 | 
33 | # 훈련이 끝나고, 훈련된 최종 adapter와 merge를 시도합니다.
34 | python merge.py \
35 |     --base_model_name_or_path $BASIS \
36 |     --peft_model_path $output_directory \
37 |     --output_dir $final_dir
38 | 
39 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | <p align="center" width="100%">
 2 | </p>
 3 | 
 4 | # Easy DPO
 5 | - DPO 코드를 한국어로 조금 더 편하고 편리하게 정할 수 있는 공유용 코드 레포입니다.
 6 | 
 7 | - 요즘 LLM은 대부분 SFT 다음에 DPO로 allign을 잡아주게 됩니다.
 8 | 
 9 | - trl 라이브러리의 기본 DPO예제를 더 편하게 이해하고 적용할 수 있도록 코드를 구성하였습니다.
10 | 
11 | - 오픈소스 커뮤니티와 비기너들에게 유익한 정보를 공유할 수 있는 **Markr.AI팀**이 될 수 있도록 더 노력하겠습니다.
12 |   
13 | # DATASET 💾
14 | ## 훈련데이터셋 📚
15 | 
16 | - 훈련 Dataset은 기본적으로 두개의 pair가 구성이된 답변이 존재해야, 적절하게 훈련이 가능합니다.
17 | 
18 | - 다음 표와 같이 데이터를 준비하시면 됩니다.
19 |   
20 | 칼럼명: | question | response_j | response_k
21 | -- | -- | -- | --
22 | **설명:** | **isntruction** | **선호답변** | **비선호 답변**
23 | 
24 | 
25 | 
26 | # Train
27 | 
28 | - 훈련 코드는 매우 간단하게 `sh`파일로 훈련이 가능하게 구성해 두었습니다. 적절한 인자를 넣어서 활용하시길 바랍니다.
29 | 
30 | - sh 파일을 구성하시고, sh 파일을 구동하시면, 간편하게 훈련을 시킬 수 있도록 파일을 구성해 두었습니다.
31 | 
32 | ```python
33 | sh DPO_train.sh
34 | ```
35 | 
36 | - sh 파일에 적절한 인자를 넣어주는게 중요합니다.
37 | 
38 | - 라이브러리 버전은 차후 제가 확인 후 다시 한번 공유 드릴 수 있도록 하겠습니다.
39 | 
40 | 
41 | 
42 | # Acknowledgement
43 | 
44 | - 주식회사 마커의 LLM프로젝트를 학술적으로 연구하였으며, 학술적인 목적임을 알리는 바입니다.
45 | 


--------------------------------------------------------------------------------
/merge.py:
--------------------------------------------------------------------------------
 1 | from transformers import AutoModelForCausalLM, AutoTokenizer
 2 | from peft import PeftModel
 3 | import torch
 4 | from huggingface_hub import login
 5 | import os
 6 | import argparse
 7 | 
 8 | def get_args():
 9 |     parser = argparse.ArgumentParser()
10 |     parser.add_argument("--base_model_name_or_path", type=str)
11 |     parser.add_argument("--token_id", type=str)
12 |     parser.add_argument("--peft_model_path", type=str)
13 |     parser.add_argument("--output_dir", type=str)
14 |     parser.add_argument("--device", type=str, default="auto")
15 | 
16 |     return parser.parse_args()
17 | 
18 | def main():
19 |     args = get_args()
20 |     login(token=args.token_id)
21 | 
22 |     if args.device == 'auto':
23 |         device_arg = { 'device_map': 'auto' }
24 |     else:
25 |         device_arg = { 'device_map': { "": args.device} }
26 | 
27 |     print(f"Loading base model: {args.base_model_name_or_path}")
28 |     base_model = AutoModelForCausalLM.from_pretrained(
29 |         args.base_model_name_or_path,
30 |         load_in_8bit=True,
31 |         return_dict=True,
32 |         torch_dtype=torch.float16,
33 |         **device_arg
34 |     )
35 |     tokenizer = AutoTokenizer.from_pretrained(args.base_model_name_or_path)
36 |     
37 |     # 만약 토크나이저가 확장되도록 변한다면 바꾸어 주면 됨.
38 |     # base_model.resize_token_embeddings(len(tokenizer))
39 | 
40 |     
41 |     print(f"Loading PEFT: {args.peft_model_path}")
42 |     model = PeftModel.from_pretrained(base_model, args.peft_model_path, **device_arg)
43 |     print(f"Running merge_and_unload")
44 |     model = model.merge_and_unload()
45 | 
46 |     
47 |     model.save_pretrained(f"{args.output_dir}")
48 |     tokenizer.save_pretrained(f"{args.output_dir}")
49 |     print(f"Model saved to {args.output_dir}")
50 | 
51 | if __name__ == "__main__" :
52 |     main()


--------------------------------------------------------------------------------
/training_DPO.py:
--------------------------------------------------------------------------------
  1 | """
  2 | @ Author: DopeorNope Lee(Seungyoo Lee)
  3 | @ Description: training LLM using the DPO Method
  4 | @ Date: 07.Aug.2023
  5 | @ Last Update: 21.Jun.2024
  6 | 
  7 | @ Notice:
  8 | Unauthorized copying and sharing of code without the author's permission are prohibited.
  9 | 
 10 | """
 11 | 
 12 | import os
 13 | from dataclasses import dataclass, field
 14 | from typing import Dict, Optional
 15 | import torch
 16 | from datasets import Dataset, load_dataset
 17 | from peft import LoraConfig
 18 | from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, TrainingArguments, BitsAndBytesConfig
 19 | from accelerate import dispatch_model, infer_auto_device_map
 20 | from accelerate.utils import get_balanced_memory
 21 | from trl import DPOTrainer
 22 | from huggingface_hub import login
 23 | import argparse
 24 | 
 25 | 
 26 | 
 27 | 
 28 | 
 29 | def get_args():
 30 | 
 31 |     """
 32 |     이 함수는 각 DPO로 훈련시키기 위해 중요한 인수인자를 sh 파일로 부터 받아와서 
 33 | 
 34 |     """
 35 | 
 36 |     parser = argparse.ArgumentParser()
 37 |     
 38 |     parser.add_argument("--num_epochs", type=int, default=1)
 39 |     parser.add_argument("--beta", type=float, default=0.1)
 40 |     parser.add_argument("--datapath", type=str, default="")
 41 |     parser.add_argument("--model_name_or_path", type=str, default="")
 42 |     parser.add_argument("--learning_rate", type=float, default=1e-5)
 43 |     parser.add_argument("--lr_scheduler_type", type=str, default="linear")
 44 |     parser.add_argument("--warmup_steps",type=int,default=100)
 45 |     parser.add_argument("--token_id", type =str, default='',help='please enter yout huggingface token ID')
 46 |     parser.add_argument("--weight_decay",type=float,default=0.05)
 47 |     parser.add_argument("--optimizer_type", type=str, default="paged_adamw_32bit")
 48 |     parser.add_argument("--per_device_train_batch_size", type=int, default=32)
 49 |     parser.add_argument("--per_device_eval_batch_size", type=int, default=1)
 50 |     parser.add_argument("--gradient_accumulation_steps", type=int, default=4)
 51 |     parser.add_argument("--gradient_checkpointing", type=bool, default=True)
 52 |     parser.add_argument("--lora_alpha", type=float, default=16)
 53 |     parser.add_argument("--lora_dropout", type=float, default=0.05)
 54 |     parser.add_argument("--lora_r", type=int, default=16)
 55 |     parser.add_argument("--max_prompt_length", type=int, default=4096)
 56 |     parser.add_argument("--max_length", type=int, default=4096)
 57 |     parser.add_argument("--max_step", type=int, default=1000)
 58 |     parser.add_argument("--logging_steps", type=int, default=10)
 59 |     parser.add_argument("--save_steps", type=int, default=100)
 60 |     parser.add_argument("--eval_steps", type=int, default=100)
 61 |     parser.add_argument("--output_dir", type=str, default="./results")
 62 |     parser.add_argument("--log_freq", type=int, default=1)
 63 |     parser.add_argument("--sanity_check", type=bool, default=False)
 64 |     parser.add_argument("--report_to", type=str, default="wandb")
 65 |     parser.add_argument("--ignore_bias_buffers", type=bool, default=False)
 66 |     parser.add_argument("--lora_target_modules",type=list, default = ['embed_tokens', 'q_proj', 'k_proj', 'v_proj', 'gate_proj', 'down_proj', 'up_proj', 'lm_head'] )
 67 |     
 68 |     return parser.parse_args()
 69 | 
 70 | 
 71 | # pair 한 데이터셋을 준비하는 함수입니다.
 72 | def paired_data_preparation(
 73 |     data_dir: str = "", #default
 74 |     sanity_check: bool = False,
 75 |     cache_dir: str = None,
 76 |     split_criteria: str = "train",
 77 |     num_proc: int=24,
 78 | ) -> Dataset:
 79 |     """
 80 |     이 데이터셋은 이후 딕셔너리 형태로 변환되며  다음과 같은 형태로 prompt, chosen, reject로 담기게 됩니다.
 81 |     {
 82 |         'prompt': List[str],
 83 |         'chosen': List[str],
 84 |         'rejected': List[str],
 85 |     }
 86 | 
 87 |     Prompt의 구조는 다음과 같이 담기게 됩니다(알파카 프롬프트):  
 88 |       "###질문: " + <prompt> + "\n\n###답변: "
 89 |     """
 90 | 
 91 |     dataset = load_dataset(data_dir, split=split_criteria ,cache_dir=cache_dir)
 92 |     
 93 |     
 94 |     original_columns = dataset.column_names
 95 | 
 96 |     if sanity_check:
 97 |         dataset = dataset.select(range(min(len(dataset), 1000)))
 98 | 
 99 |     # 함수내에 정의되는 함수로, 각 데이터 포맷과 구조에 맞게 매핑해주는 역할을 진행합니다.
100 |     def return_prompt_and_responses(samples) -> Dict[str, str]:
101 |         return {
102 |             "prompt": ["###질문:\n" + question + "\n\n###답변:\n" for question in samples["question"]],
103 |             "chosen": samples["response_j"],
104 |             "rejected": samples["response_k"],
105 |         }
106 | 
107 |     return dataset.map(
108 |         return_prompt_and_responses,
109 |         batched=True,
110 |         num_proc=num_proc,
111 |         remove_columns=original_columns,
112 |     )
113 | 
114 | 
115 | 
116 | 
117 | 
118 | def main():
119 |     args = get_args()
120 | 
121 | 
122 |     print(
123 |             f"################################ Training Arguments ################################"
124 |             f"model_name_or_path: {args.model_name_or_path}\n"
125 |             f"datapath: {args.datapath}\n"
126 |             f"output_dir: {args.output_dir}\n"
127 |             f"per_device_train_batch_size: {args.per_device_train_batch_size}\n"
128 |             f"per_device_eval_batch_size: {args.per_device_eval_batch_size}\n"
129 |             f"num_epochs: {args.num_epochs}\n"
130 |             f"max_step: {args.max_step}\n"
131 |             f"learning_rate: {args.learning_rate}\n"
132 |             f"cutoff_len(max_length): {args.max_length}\n"
133 |             f"lora_r: {args.lora_r}\n"
134 |             f"lora_alpha: {args.lora_alpha}\n"
135 |             f"lora_dropout: {args.lora_dropout}\n"
136 |             f"lora_target_modules: {args.lora_target_modules}\n"
137 |             f"max_prompt_length: {args.max_prompt_length}\n"
138 |             f"##################################################################################\n"
139 |             
140 |     )
141 | 
142 |     login(token=args.token_id)
143 | 
144 | 
145 |     # SFT MODEL
146 |     model = AutoModelForCausalLM.from_pretrained(
147 |         args.model_name_or_path,
148 |         low_cpu_mem_usage=True,
149 |         torch_dtype=torch.float16,
150 |         trust_remote_code=True,      
151 |         device_map="auto")
152 | 
153 | 
154 |     model.config.use_cache = False
155 |     model.is_parallelizable = True
156 |     model.model_parallel = True
157 | 
158 |     model.config.max_position_embeddings= args.max_prompt_length
159 | 
160 |     print("model's max_position_embeddings :",model.config.max_position_embeddings)
161 | 
162 | 
163 |     if args.ignore_bias_buffers:
164 |         # torch distributed hack
165 |         model._ddp_params_and_buffers_to_ignore = [
166 |             name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool
167 |         ]
168 | 
169 |     
170 |     tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
171 |     tokenizer.pad_token = tokenizer.eos_token
172 | 
173 |     #model.resize_token_embeddings(len(tokenizer))
174 |     
175 |     train_dataset = paired_data_preparation(data_dir= args.datapath, split_criteria= "train", sanity_check=args.sanity_check)
176 |     
177 |     train_dataset = train_dataset.filter(
178 |     
179 |         lambda x: len(x["prompt"]) + len(x["chosen"]) <= args.max_length
180 |     
181 |         and len(x["prompt"]) + len(x["rejected"]) <= args.max_length)
182 |     
183 | 
184 |     eval_dataset = paired_data_preparation(data_dir= args.datapath, split_criteria= "validation", sanity_check=True)
185 |     
186 |     eval_dataset = eval_dataset.filter(
187 |     
188 |        lambda x: len(x["prompt"]) + len(x["chosen"]) <= args.max_length
189 |     
190 |        and len(x["prompt"]) + len(x["rejected"]) <= args.max_length)
191 |     
192 | 
193 | 
194 |     training_args = TrainingArguments(
195 |         num_train_epochs= args.num_epochs,
196 |         per_device_train_batch_size = args.per_device_train_batch_size,
197 |         per_device_eval_batch_size = args.per_device_eval_batch_size,
198 |        # max_steps = args.max_step,
199 |         logging_steps = args.logging_steps,
200 |         save_steps = args.save_steps,
201 |         gradient_accumulation_steps = args.gradient_accumulation_steps,
202 |         gradient_checkpointing = args.gradient_checkpointing,
203 |         learning_rate = args.learning_rate,
204 |         evaluation_strategy="steps",
205 |         eval_steps=args.eval_steps,
206 |         output_dir = args.output_dir,
207 |         #report_to = args.report_to,
208 |         lr_scheduler_type = args.lr_scheduler_type,
209 |         warmup_steps = args.warmup_steps,
210 |         optim = args.optimizer_type,
211 |         bf16 = True,
212 |         remove_unused_columns = False,
213 |         run_name = "dpo_llama2",
214 |     )
215 | 
216 | 
217 |     peft_config = LoraConfig(
218 |         r = args.lora_r,
219 |         lora_alpha = args.lora_alpha,
220 |         lora_dropout = args.lora_dropout,
221 |         target_modules = args.lora_target_modules,
222 |         bias = "none",
223 |         task_type = "CAUSAL_LM")
224 |     
225 |         
226 |     print("###################################################################################")    
227 |     print("############################  MODEL was Lodaded in GPU ############################")
228 |     print("###################################################################################")    
229 | 
230 |     dpo_trainer = DPOTrainer(
231 |         model,
232 |         ref_model = None,   # ref 모델을 None으로 놓게 되면 SFT + adapter가 붙은 모델에서 adapter를 떼고, policy에 따른 최적화를 진행하게 됩니다. 두개의 모델을 로드할 필요가 없어 메모리 이득을 꾀할 수 있습니다.
233 |         args = training_args,
234 |         beta = args.beta,
235 |         train_dataset= train_dataset,
236 |         eval_dataset = eval_dataset,
237 |         tokenizer = tokenizer,
238 |         peft_config = peft_config,
239 |         max_prompt_length = args.max_prompt_length,
240 |         max_length = args.max_length,
241 |     )
242 | 
243 |     print("###################################################################################")    
244 |     print("########################  Trainin Process is preparing now  #######################")
245 |     print("###################################################################################")    
246 | 
247 |     dpo_trainer.train()
248 |     dpo_trainer.save_model(args.output_dir)
249 | 
250 |     output_dir = os.path.join(args.output_dir, "final_checkpoint")
251 |     dpo_trainer.model.save_pretrained(output_dir)
252 | 
253 | 
254 | if __name__ == "__main__" :
255 |     main()


--------------------------------------------------------------------------------