Trainer 메서드 서브클래싱

Trainer 메서드를 서브클래스화하여 전체 루프를 다시 작성하지 않고도 학습 동작을 변경할 수 있습니다. 서브클래스화는 training loop를 수정하는 것으로, 예를 들어 forward pass나 loss 계산 방식을 바꿀 수 있습니다.

서브클래스를 만들기 전에, Trainer가 무엇을 계산하는지 또는 언제 그리고 어떤 조건에서 동작하는지를 변경해야 하는지 고려해 보세요. 타이밍 및 조건부 로직이 필요할 경우에는 대신 Callback을 사용하세요. Callback은 이벤트가 발생하는 시점(로그 기록, 평가, 조기 종료)을 제어하고, 서브클래스화는 이벤트의 내용(손실 계산, 데이터 로딩, 최적화)을 변경합니다.

자식 클래스로 만들 수 있는 메서드의 전체 목록은 Trainer API 문서를 참조하세요. _로 시작하는 비공개 메서드, 예를 들어 _save_checkpoint나 _evaluate도 오버라이드할 수 있지만, 이러한 메서드는 사전 통보 없이 변경될 수 있습니다.

get_train_dataloader

표준 get_train_dataloader() 메서드는 한 배치(batch)를 로드하고, 해당 배치로 학습한 후 버리고, 다음 배치를 로드합니다.

def get_train_dataloader(self):
    return self._get_dataloader(
        batch_size=self._train_batch_size,
        ...
)

GRPO는 학습 전에 completions을 생성하는 온라인 강화 학습 알고리즘입니다. 매 스텝마다 completions을 생성하는 것은 자기회귀적(autoregressive) 특성 때문에 비용이 많이 듭니다. 512-token의 completion은 학습 스텝 하나당 1회의 forward pass에 비해 약 512회의 순차적 forward pass가 필요합니다. GRPOTrainer는 get_train_dataloader()를 서브클래스화하여 여러 스텝에 걸쳐 생성을 배치 처리합니다.

trl.GRPOTrainer.get_train_dataloader는 steps_per_generation 인수를 배치 크기에 곱하여 여러 학습 단계에 사용할 배치를 한 번에 로드합니다. train_batch_size=4이고 steps_per_generation=8인 경우, 데이터 로더는 크기가 32인 배치를 생성하여 생성 비용을 8배 줄입니다.

def get_train_dataloader(self):
    dataloader_params = {
        "batch_size": self._train_batch_size * self.args.steps_per_generation, # this is the only change
        ...
    }

손실 계산

compute_loss()은 모델이 계산한 교차 엔트로피 손실을 반환합니다.

def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
    ...
    outputs = model(**inputs)
    ...
    loss = outputs["loss"] # get loss from model

    return (loss, outputs) if return_outputs else loss

DPO는 정책 모델이 거부된 응답에 비해 선택된 응답을 얼마나 선호하는지를 참조 모델에 대해 측정합니다. DPOTrainer는 손실 계산이 표준 교차 엔트로피와 여러 면에서 다르기 때문에 compute_loss()를 재정의합니다:

모델은 라벨을 절대 확인하지 않으며, DPO가 log-probs를 계산할 수 있도록 logits만 반환합니다
chosen 및 rejected 응답이 연결됩니다
reference model은 자체 log-probs를 계산합니다
손실은 π_chosen, π_rejected, π_ref_chosen, π_ref_rejected의 함수입니다

위 방법 중 어느 것도 표준 Trainer.compute_loss() 메서드에 해당하지 않습니다.

def compute_loss(
    self,
    model: PreTrainedModel | nn.Module,
    inputs: dict[str, torch.Tensor | Any],
    return_outputs=False,
    num_items_in_batch=None,
) -> torch.Tensor | tuple[torch.Tensor, dict[str, float]]:
    ...
    outputs = model(**inputs)
    logits = outputs.logits
    logps = get_logps(logits, inputs)
    chosen_logps, rejected_logps = logps.chunk(2, dim=0)  # batch is [chosen, rejected]
    ref_logits = self.ref_model(**inputs).logits
    ref_logps = get_logps(ref_logits, inputs)
    ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0)  # batch is [chosen, rejected]
    chosen_scores = chosen_logps - ref_chosen_logps
    rejected_scores = rejected_logps - ref_rejected_logps
    per_sequence_loss = -F.logsigmoid(self.beta * chosen_scores - rejected_scores)
    loss = per_sequence_loss.mean()
    return (loss, outputs) if return_outputs else loss

다음 단계

실제 사용 사례를 더 확인하려면, TRL에서 GRPOTrainer와 DPOTrainer가 Trainer를 어떻게 확장하는지, 또는 Axolotl가 이를 기반으로 커스텀 트레이너를 어떻게 구축하는지 살펴보세요.
트레이닝 이벤트(예: 트레이닝 스텝 종료 시 메트릭 로깅)에서 발생하는 동작만 커스터마이즈해야 하는 경우 Callbacks 가이드를 확인하세요.

Update on GitHub

Transformers

Trainer 메서드 서브클래싱

get_train_dataloader

손실 계산

다음 단계