Hugging Face Transformers 팀이 그래디언트 누적(Gradient Accumulation) 중 손실 계산 방식을 수정해 토큰 레벨 작업에서 수학적 동등성 보장

Context

Transformers 라이브러리의 기본 손실 함수는 배치 단위로 평균 손실을 계산하도록 설계되었다. 그래디언트 누적을 활성화했을 때 비활성화 상태와 손실값이 일치하지 않는 문제가 발생했다. 이는 토큰 분류나 인과 언어 모델(Causal LM) 같은 토큰 레벨 작업에서 올바른 손실 계산 방식과 기본 구현의 불일치에서 비롯되었다.

Technical Solution

• 그래디언트 누적 중 손실 계산 방식 변경: 누적된 모든 배치의 총 손실을 해당 배치들의 전체 패딩되지 않은 토큰 개수로 나누기(기존: 배치 단위 손실의 평균)
• 손실 함수 내부에 num_items 파라미터 추가: CrossEntropy의 reduction="sum" 방식으로 변경해 누적 단계에서 올바른 정규화 수행
• PreTrainedModel에 loss_function 속성 추가: 사용자가 config.loss_type을 통해 커스텀 손실 함수 정의 가능하도록 LOSS_MAPPING 제공
• Trainer API 확장: 사용자가 자체 손실 함수를 Trainer에 직접 전달할 수 있도록 지원 (PR #34198)
• 점진적 모델 지원: PR #34191에서 인기 있는 모델부터 수정 시작, 차후 기여자 모집을 통해 전체 모델로 확대

Key Takeaway

기본 API가 직관적이지 않으면 예상치 못한 버그의 원인이 된다. 특정 사용 사례(그래디언트 누적)에서만 필요한 동작을 명시적으로 지원하고, 동시에 고급 사용자를 위한 커스터마이징 경로를 함께 제공해야 한다.