MetaAI의 Self-Rewarding Language Model에서 제안된 훈련 프레임워크 구현
그들은 DPO 논문의 제목을 마음에 새겼습니다.
이 라이브러리에는 Nous Research의 Teknium이 낙관론을 표명한 SPIN 구현도 포함되어 있습니다.
$ pip install self-rewarding-lm-pytorch
import torch
from torch import Tensor
from self_rewarding_lm_pytorch import (
SelfRewardingTrainer ,
create_mock_dataset
)
from x_transformers import TransformerWrapper , Decoder
transformer = TransformerWrapper (
num_tokens = 256 ,
max_seq_len = 1024 ,
attn_layers = Decoder (
dim = 512 ,
depth = 1 ,
heads = 8
)
)
sft_dataset = create_mock_dataset ( 100 , lambda : ( torch . randint ( 0 , 256 , ( 256 ,)), torch . tensor ( 1 )))
prompt_dataset = create_mock_dataset ( 100 , lambda : 'mock prompt' )
def decode_tokens ( tokens : Tensor ) -> str :
decode_token = lambda token : str ( chr ( max ( 32 , token )))
return '' . join ( list ( map ( decode_token , tokens )))
def encode_str ( seq_str : str ) -> Tensor :
return Tensor ( list ( map ( ord , seq_str )))
trainer = SelfRewardingTrainer (
transformer ,
finetune_configs = dict (
train_sft_dataset = sft_dataset ,
self_reward_prompt_dataset = prompt_dataset ,
dpo_num_train_steps = 1000
),
tokenizer_decode = decode_tokens ,
tokenizer_encode = encode_str ,
accelerate_kwargs = dict (
cpu = True
)
)
trainer ( overwrite_checkpoints = True )
# checkpoints after each finetuning stage will be saved to ./checkpoints
SPIN은 다음과 같이 교육할 수 있습니다. Readme의 마지막 예에 표시된 대로 미세 조정 파이프라인에 추가할 수도 있습니다.
import torch
from self_rewarding_lm_pytorch import (
SPINTrainer ,
create_mock_dataset
)
from x_transformers import TransformerWrapper , Decoder
transformer = TransformerWrapper (
num_tokens = 256 ,
max_seq_len = 1024 ,
attn_layers = Decoder (
dim = 512 ,
depth = 6 ,
heads = 8
)
)
sft_dataset = create_mock_dataset ( 100 , lambda : ( torch . randint ( 0 , 256 , ( 256 ,)), torch . tensor ( 1 )))
spin_trainer = SPINTrainer (
transformer ,
max_seq_len = 16 ,
train_sft_dataset = sft_dataset ,
checkpoint_every = 100 ,
spin_kwargs = dict (
λ = 0.1 ,
),
)
spin_trainer ()
LLM 판사가 아닌 자신만의 보상 프롬프트를 실험하고 싶다고 가정해 보겠습니다. 먼저 RewardConfig
가져와야 합니다. 다음으로 이를 reward_prompt_config
로 트레이너에 전달해야 합니다.
# first import
from self_rewarding_lm_pytorch import RewardConfig
# then say you want to try asking the transformer nicely
# reward_regex_template is the string that will be looked for in the LLM response, for parsing out the reward where {{ reward }} is defined as a number
trainer = SelfRewardingTrainer (
transformer ,
...,
self_reward_prompt_config = RewardConfig (
prompt_template = """
Pretty please rate the following user prompt and response
User: {{ prompt }}
Response: {{ response }}
Format your score as follows:
Rating: <rating as integer from 0 - 10>
""" ,
reward_regex_template = """
Rating: {{ reward }}
"""
)
)
마지막으로, 임의의 순서로 미세 조정을 실험하려는 경우 FinetuneConfig
인스턴스를 목록으로 finetune_configs
에 전달하면 유연성도 확보할 수 있습니다.
전. 인터리빙 SPIN, 외부 보상, 자기 보상에 대한 연구를 수행하고 싶다고 가정해 보겠습니다.
이 아이디어는 개인 디스코드 채널의 Teknium에서 시작되었습니다.
# import the configs
from self_rewarding_lm_pytorch import (
SFTConfig ,
SelfRewardDPOConfig ,
ExternalRewardDPOConfig ,
SelfPlayConfig ,
)
trainer = SelfRewardingTrainer (
model ,
finetune_configs = [
SFTConfig (...),
SelfPlayConfig (...),
ExternalRewardDPOConfig (...),
SelfRewardDPOConfig (...),
SelfPlayConfig (...),
SelfRewardDPOConfig (...)
],
...
)
trainer ()
# checkpoints after each finetuning stage will be saved to ./checkpoints
일괄 처리의 다양한 위치에서 진행될 수 있도록 샘플링을 일반화하고 모든 샘플링을 일괄 처리하도록 수정합니다. 또한 일부 사람들이 이를 허용하는 상대 위치를 가진 변환기를 가지고 있는 경우 왼쪽 패딩 시퀀스를 허용합니다.
EOS를 처리하다
기본 llm-as-judge 대신 자신만의 보상 프롬프트를 사용하는 예를 보여주세요.
쌍을 샘플링하기 위한 다양한 전략을 허용합니다.
조기 스토퍼
sft, spin, self-rewarding dpo, 외부 보상 모델이 있는 dpo 순서
보상에 대한 검증 기능을 허용합니다(보상은 정수, 부동 소수점, 특정 범위 사이 등이어야 함).
kv 캐시의 다양한 구현을 가장 잘 처리하는 방법을 알아내세요. 지금은 그냥 사용하지 마세요.
모든 체크포인트 폴더를 자동으로 지우는 환경 플래그
@misc { yuan2024selfrewarding ,
title = { Self-Rewarding Language Models } ,
author = { Weizhe Yuan and Richard Yuanzhe Pang and Kyunghyun Cho and Sainbayar Sukhbaatar and Jing Xu and Jason Weston } ,
year = { 2024 } ,
eprint = { 2401.10020 } ,
archivePrefix = { arXiv } ,
primaryClass = { cs.CL }
}
@article { Chen2024SelfPlayFC ,
title = { Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models } ,
author = { Zixiang Chen and Yihe Deng and Huizhuo Yuan and Kaixuan Ji and Quanquan Gu } ,
journal = { ArXiv } ,
year = { 2024 } ,
volume = { abs/2401.01335 } ,
url = { https://api.semanticscholar.org/CorpusID:266725672 }
}
@article { Rafailov2023DirectPO ,
title = { Direct Preference Optimization: Your Language Model is Secretly a Reward Model } ,
author = { Rafael Rafailov and Archit Sharma and Eric Mitchell and Stefano Ermon and Christopher D. Manning and Chelsea Finn } ,
journal = { ArXiv } ,
year = { 2023 } ,
volume = { abs/2305.18290 } ,
url = { https://api.semanticscholar.org/CorpusID:258959321 }
}
@inproceedings { Guo2024DirectLM ,
title = { Direct Language Model Alignment from Online AI Feedback } ,
author = { Shangmin Guo and Biao Zhang and Tianlin Liu and Tianqi Liu and Misha Khalman and Felipe Llinares and Alexandre Rame and Thomas Mesnard and Yao Zhao and Bilal Piot and Johan Ferret and Mathieu Blondel } ,
year = { 2024 } ,
url = { https://api.semanticscholar.org/CorpusID:267522951 }
}