Hugging Face Transformers Gradient Accumulation 수정
Hugging Face는 transformers Trainer에서 gradient accumulation이 전체 배치 학습과 동등한 결과를 내지 못하는 수학적 불일치를 수정했습니다. 이번 수정으로 손실은 배치당 손실 값을 단순 평균하는 것이 아니라, accumulation 단계의 모든 배치에 걸친 비패딩 토큰 총 수를 기준으로 계산됩니다.
Gradient Accumulation 오류의 근본 원인
Gradient accumulation은 더 큰 배치 크기로 학습하는 것과 수학적으로 동일하도록 설계되었습니다. 그러나 gradient accumulation을 켜거나 끄는 경우 손실이 일치하지 않는 불일치가 발견되었습니다.
이 문제는 transformers 모델의 모델링 코드 내에서 제공되는 "default" 손실 함수에서 비롯되었습니다. 이러한 함수는 labels와 input_ids가 모델에 전달될 때 자동으로 호출되어 사용자의 API 사용을 단순화합니다. Causal Language Modeling (Causal LM)과 같은 토큰 수준 작업의 경우, 올바른 손실 계산은 gradient accumulation 단계의 모든 배치에 대한 총 손실을 해당 배치들의 비패딩 토큰 총 수로 나누어야 합니다. 이전 구현은 배치별 손실 값의 평균을 계산했으며, 이는 해당 사용 사례에 대해 수학적으로 올바르지 않았습니다.
수정의 기술적 구현
정밀도 및 계산 오류를 해결하기 위해 Hugging Face는 손실 계산 로직을 수정했습니다. 핵심 변경 사항은 기본 평균 감소(mean reduction)에서 합계 감소(sum reduction)로 전환한 뒤, 전체 항목 수(num_items)로 나누는 것입니다:
# Corrected loss calculation
loss = nn.functional.cross_entropy(shift_logits, shift_labels, ignore_index=-100, reduction="sum")
loss = loss / num_items
장기적인 아키텍처 변경 사항
Hugging Face는 향후 손실 계산 문제를 방지하고 유연성을 높이기 위해 두 가지 주요 구조적 변경을 구현하고 있습니다:
1. 자동 손실 보정
기본 손실 함수를 사용하는 사용자에 대해, 라이브러리는 gradient accumulation 동안 필요한 변경을 자동으로 적용하여 보고 및 사용되는 손실이 정확하도록 합니다.
2. 사용자 정의 가능한 손실 API
내부 라이브러리 문제로 사용자가 차단되는 것을 방지하기 위해, Hugging Face는 사용자가 자신의 손실 함수를 직접 Trainer에 전달할 수 있는 API를 도입하고 있습니다.
모든 PreTrainedModel을 상속하는 모델은 이제 loss_function 속성을 포함합니다. 이 속성은 config.loss_type에 의해 결정되며, 사용자는 LOSS_MAPPING을 수정하여 손실을 커스터마이즈할 수 있습니다:
def my_super_loss(logits, labels):
return nn.functional.cross_entropy(logits, labels, ignore_index=-100)
LOSS_MAPPING["my_loss_type"] = my_super_loss
배포 및 사용 가능성
이번 수정은 두 개의 특정 Pull Request를 통해 롤아웃되고 있습니다:
- PR #34191: 가장 인기 있는 모델에 대한 첫 번째 변경을 구현합니다.
- PR #34198: 두 번째 변경을 구현하여 사용자가 자체 손실 함수를 제공하고 배치당 관측된 샘플 수를 계산에 활용할 수 있게 합니다.
사용자는 main 브랜치에서 transformers 라이브러리를 설치하여 즉시 이 수정들을 적용할 수 있습니다:
pip install git+https://github.com/huggingface/transformers
Sources
- OriginalFixing Gradient Accumulation