Hugging Face's logo

TRL documentation

참조 모델용 BEMA

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.17.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

참조 모델용 BEMA

이 기능은 DPO 학습 중에 참조 모델을 업데이트하기 위해 BEMA 알고리즘을 구현합니다.

사용

from trl.experimental.bema_for_ref_model import BEMACallback, DPOTrainer
from datasets import load_dataset

dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train")

bema_callback = BEMACallback(update_ref_model=True)

trainer = DPOTrainer(
    model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
    train_dataset=dataset,
    callbacks=[bema_callback],
)
trainer.train()

DPOTrainer

클래스 trl.DPOTrainer

< >

( *args**kwargs )

train

< >

( resume_from_checkpoint: str | bool | None = Nonetrial: optuna.Trial | dict[str, Any] | None = Noneignore_keys_for_eval: list[str] | None = None ) → ~trainer_utils.TrainOutput

매개변수

  • resume_from_checkpoint (str 또는 bool, 선택 사항) — str인 경우, 이전 Trainer 인스턴스가 저장한 체크포인트의 로컬 경로입니다. bool이고 True와 같으면, 이전 Trainer 인스턴스가 args.output_dir에 저장한 마지막 체크포인트를 로드합니다. 이 값이 지정되어 있으면, 여기서 로드된 모델/옵티마이저/스케줄러 상태로부터 학습이 재개됩니다.
  • trial (optuna.Trial 또는 dict[str, Any], 선택 사항) — 하이퍼파라미터 탐색을 위한 트라이얼 실행 또는 하이퍼파라미터 사전입니다.
  • ignore_keys_for_eval (list[str], 선택 사항) — 모델의 출력(사전에 해당하는 경우)에서 학습 중 평가 예측을 수집할 때 무시해야 할 키의 목록입니다.

반환값

~trainer_utils.TrainOutput

전체 스텝 수, 학습 손실, 메트릭을 포함하는 객체.

주요 학습 진입점.

save_model

< >

( output_dir: str | None = None_internal_call: bool = False )

모델을 저장하므로 from_pretrained()를 사용하여 다시 로드할 수 있습니다.

main 프로세스에서만 저장합니다.

push_to_hub

< >

( commit_message: str | None = 'End of training'blocking: bool = Truetoken: str | None = Nonerevision: str | None = None**kwargs )

매개변수

  • commit_message (str, 선택 사항, 기본값 "End of training") — 푸시할 때 커밋할 메시지입니다.
  • blocking (bool, 선택 사항, 기본값 True) — git push가 완료된 후에만 함수가 반환되어야 하는지 여부.
  • token (str, 선택 사항, 기본값 None) — Trainer의 원래 args를 덮어쓸 수 있는 쓰기 권한이 있는 토큰.
  • revision (str, 선택 사항) — 커밋할 git revision입니다. 기본값은 “main” 브랜치의 head입니다.
  • kwargs (dict[str, Any], 선택) — ~Trainer.create_model_card에 함께 전달되는 추가 키워드 인자입니다.

self.args.hub_model_id 저장소에 self.model과 self.processing_class를 🤗 모델 허브에 업로드합니다.

BEMACallback

클래스 trl.BEMACallback

< >

( update_freq: int = 400ema_power: float = 0.5bias_power: float = 0.2lag: int = 10update_after: int = 0multiplier: float = 1.0min_ema_multiplier: float = 0.0device: str = 'cpu'update_ref_model: bool = Falseref_model_update_freq: int = 400ref_model_update_after: int = 0 )

매개변수

  • update_freq (int, 선택 사항, 기본값은 400) — X단계마다 BEMA 가중치를 업데이트합니다. 논문에서는 이를 {@html "ϕ \\phi ϕ"} in the paper.
  • ema_power (float, optional, defaults to 0.5) — Power for the EMA decay factor. Denoted {@html "κ \\kappa κ "} 로 표기합니다. EMA를 비활성화하려면 이 값을 0.0.
  • bias_power (float, 선택 사항, 기본값은 0.2) — BEMA 스케일링 인자의 지수입니다. 논문에서는 {@html "η \\eta η"}로 표기합니다. 값이 크면(예: 8.0) {@html "αt \\alpha_t αt​"}이 0으로 감소하여 bias-correction 비활성화를 근사합니다. 0.0이면 모든 단계에서 {@html "αt \\alpha_t αt​"}이 1로 고정되어(최대, 감쇠되지 않은 보정) 됩니다. lag (int, 선택 사항, 기본값 10) — weight decay 스케줄의 초기 오프셋으로, 업데이트의 가상 시작 연령 역할을 하여 초기 단계의 매끄러움을 제어합니다. 논문에서는 {@html "ρ \\rho ρ "}η \\eta η"} in the paper. A large value (e.g. 8.0) makes {@html "αt \\alpha_t αt​"} decay to 0, approximating disabled bias-correction; 0.0 instead pins {@html "αt \\alpha_t αt​"} at 1 for every step (maximum, undecayed correction).
  • lag (int, optional, defaults to 10) — Initial offset in the weight decay schedule that controls early-stage smoothness by acting as a virtual starting age for the updates. Denoted as {@html "ρ \\rho ρ"} 로 표기합니다.
  • update_after (int, 선택 사항, 기본값은 0) — BEMA 가중치를 업데이트하기 시작하기 전의 버닝인 시간입니다. 논문에서는 {@html "τ \\tau τ"}로 표기합니다. multiplier (float, 선택 사항, 기본값 1.0) — EMA 감쇠 계수의 초기 값입니다. 논문에서는 {@html "γ \\gamma γ "}로 표기합니다.τ \\tau τ"} in the paper.
  • multiplier (float, optional, defaults to 1.0) — Initial value for the EMA decay factor. Denoted as {@html "γ \\gamma γ"} 입니다.
  • min_ema_multiplier (float, 선택 사항, 기본값 0.0) — EMA 감쇠 계수의 최소값.
  • device (str, 선택 사항, 기본값 "cpu") — BEMA 버퍼에 사용할 장치입니다. 예를 들어 "cpu" 또는 "cuda"입니다. 대부분의 경우 OOM을 방지하기 위해 이 장치는 학습에 사용되는 장치와 다르게 설정해야 합니다.
  • update_ref_model (bool, 선택 사항, 기본값 False) — BEMA 가중치로 참조 모델을 업데이트할지 여부입니다. 이를 통해 주 모델의 지연 및 평활화 버전을 참조 모델로 생성합니다.
  • ref_model_update_freq (int, 선택 사항, 기본값 400) — 이 값만큼의 스텝마다 BEMA 가중치로 참조 모델을 업데이트합니다.
  • ref_model_update_after (int, 선택 사항, 기본값 0) — 참조 모델을 업데이트하기 전에 대기할 스텝 수입니다.

Adam Block와 Cyril Zhang의 BEMA(Bias-Corrected Exponential Moving Average)를 구현한 TrainerCallback입니다. https://github.com/abblock/bema의 코드를 MIT 라이선스에 따라 사용합니다.

BEMA는 다음과 같이 스케일링되는 모델 가중치를 계산합니다: θt′=αt⋅(θt−θ0)+EMAt \theta_t' = \alpha_t \cdot (\theta_t - \theta_0) + \text{EMA}_t

여기서 θt \theta_t 은 현재 모델 가중치이고, θ0 \theta_0 은 첫 번째 update_after 단계에서의 모델 가중치 스냅샷이며, EMAt \text{EMA}_t 은 모델 가중치의 지수 이동 평균이고,αt \alpha_t 은 단계 수 t t 에 따라 감소하는 스케일링 인자입니다. αt=(ρ+γ⋅t)−η. \alpha_t = (\rho + \gamma \cdot t)^{-\eta}.

EMA는 다음과 같이 계산됩니다: EMAt=(1−βt)⋅EMAt−1+βt⋅θt \text{EMA}_t = (1 - \beta_t) \cdot \text{EMA}_{t-1} + \beta_t \cdot \theta_t

여기서 βt \beta_t 은 스텝 수에 따라 감소하는 감쇠 계수이며 t t 다음과 같습니다 βt=(ρ+γ⋅t)−κ. \beta_t = (\rho + \gamma \cdot t)^{-\kappa}.

예:

>>> from trl import BEMACallback

>>> trainer = Trainer(..., callbacks=[BEMACallback()])
Update on GitHub

←비동기 GRPO CPO→