Hugging Face 훈련 효율성: Flash Attention 2와 패킹

Hugging Face는 지시 튜닝 예제에 대한 경계 인식 패킹을 도입했으며, 이제 Flash Attention 2와 호환됩니다. 이 업데이트를 통해 패딩 없이 시퀀스를 연결할 수 있게 되어, 훈련 처리량을 최대 2배 증가시키고 피크 메모리 사용량을 20% 감소시키면서도 수렴 품질에 영향을 주지 않습니다.

향상된 훈련 처리량 및 메모리 효율성

패딩 없이 예제를 패킹하면 관련 없는 패딩 토큰으로 인한 계산 오버헤드가 크게 감소합니다. 개선 정도는 훈련 데이터셋 내 시퀀스 길이의 변동성에 따라 달라집니다:

높은 변동성 데이터셋 (예: FLAN): FLAN 데이터셋에서 훈련한 결과 Llama 2-7B, Mistral-7B, Granite-8B-code 모델에 대해 처리량이 2배 증가했으며, 피크 메모리 사용량이 20% 감소했습니다.
낮은 변동성 데이터셋 (예: OrcaMath): 예제가 더 길고 변동성이 낮은 OrcaMath 데이터셋에서 훈련한 결과 처리량이 1.4배 증가하고 피크 메모리가 6% 감소했습니다.

새 구현이 미니배치를 유지하고 패딩 훈련과 동일한 최적화 단계 수를 유지하기 때문에 검증 손실이 동일하게 유지되어 훈련 수렴에 대한 저하가 없습니다.

기술 구현: 경계 인식

이전 패킹 구현은 Flash Attention 2를 사용할 때 예제 경계를 무시하는 경우가 많아, 원치 않는 교차 예제 어텐션이 발생해 모델 품질을 저하시켰습니다. Hugging Face는 패킹 중에 경계 인식을 유지함으로써 이를 해결합니다.

이는 position_ids를 Flash Attention 2에 제공하고 flash_attn_varlen_func를 활용하여 각 미니배치에 대한 누적 시퀀스 길이(cu_seqlens)를 계산함으로써 구현됩니다. 이 방법을 통해 모델은 시퀀스를 하나의 텐서로 연결하면서도 어텐션이 올바른 시퀀스 경계에만 제한되도록 보장합니다.

지원되는 모델

이 솔루션은 모델이 position_ids를 노출해야 합니다. 현재 14개의 모델이 지원되며, 포함되는 모델은 다음과 같습니다:

Llama 2 and 3
Mistral and Mixtral
Granite
DBRX
Falcon
Gemma
OLMo
Phi 1, 2, and 3 (including phi3)
Qwen 2 and 2 MoE
StableLM
StarCoder 2

통합 및 사용법

사용자는 사용 중인 라이브러리에 따라 두 가지 주요 경로를 통해 Flash Attention 2와 함께 패킹을 구현할 수 있습니다:

Hugging Face Trainer 사용

Transformers 라이브러리의 Trainer 클래스를 사용할 때 이 기능을 활용하려면, 사용자는 다음을 수행해야 합니다:

attn_implementation="flash_attention_2" 로 모델을 인스턴스화합니다.
DataCollatorWithFlattening 콜레이터를 사용합니다.

TRL SFTTrainer 사용

TRL 라이브러리의 SFTTrainer를 DataCollatorForCompletionOnlyLM과 함께 사용하는 경우, 요구 사항은 다음과 같습니다:

Flash Attention 2를 사용하도록 모델을 인스턴스화합니다.
DataCollatorForCompletionOnlyLM을 호출할 때 padding_free=True 플래그를 설정합니다.

결론

패딩 토큰을 제거하고 경계 인식 어텐션을 구현함으로써, Hugging Face는 지시 튜닝의 효율성을 크게 향상시켰습니다. 처리량과 메모리에서 가장 큰 향상은 예제 길이의 변동성이 높은 데이터셋을 대상으로 훈련할 때 실현됩니다.

SUMMARY: Hugging Face는 지시 튜닝 예제에 대한 경계 인식 패킹을 도입했으며, Flash Attention 2와 함께 사용할 경우 훈련 처리량을 최대 2배 증가시키고 피크 메모리를 20% 감소시킵니다.

TITLE: Hugging Face 훈련 효율성: Flash Attention 2와 패킹