Flash Attention은 안정적인가요? Meta와 Harvard는 모델 중량 편차가 ​​수십 배로 변동한다는 사실을 발견했습니다.

Meta FAIR는 Harvard와 결합하여 대규모 기계 학습으로 인한 데이터 편향을 최적화하기 위한 새로운 연구 프레임워크를 제공합니다.

우리 모두 알고 있듯이 대규모 언어 모델을 훈련하는 데는 수개월이 걸리고 수백 또는 수천 개의 GPU를 사용하는 경우가 많습니다. LLaMA2 70B 모델을 예로 들면, 훈련에는 총 1,720,320 GPU 시간이 필요합니다. 대규모 모델을 교육하면 이러한 워크로드의 규모와 복잡성으로 인해 고유한 체계적 문제가 발생합니다.

최근 많은 기관에서 SOTA 생성 AI 모델을 훈련할 때 훈련 과정에서 불안정성을 보고했습니다. 이는 일반적으로 Google의 PaLM 모델과 같이 훈련 과정에서 최대 20번까지 손실이 발생하는 형태로 나타납니다. 스파이크.

수치 편차는 이러한 훈련 부정확성의 근본 원인입니다. 대규모 언어 모델 훈련의 실행 비용이 매우 높기 때문에 수치 편차를 정량화하는 방법이 중요한 문제가 되었습니다.

최신 연구에서 Meta와 Harvard University의 연구원들은 훈련 최적화의 수치 편향을 이해하기 위한 원칙적인 정량적 방법을 개발했습니다. 이는 다양한 최첨단 최적화 기술을 평가하고 대규모 모델을 훈련하는 데 사용될 때 예상치 못한 불안정성을 초래할 수 있는지 여부를 결정하는 데 사용됩니다. 연구진은 기존 최적화 방법이 일부 작업에서는 잘 수행되었지만 대형 모델에 적용하면 일부 수치 편차가 발생한다는 사실을 발견했습니다. 이러한 수치적 편향은 훈련 과정에서 불안정성을 야기하여 모델의 성능을 저하시킬 수 있습니다. 이 문제를 해결하기 위해 연구진은 원칙적인 정량적 방법을 기반으로 한 최적화를 제안했습니다.

• 논문 제목: Is Flash Attention Stable?
• 문서 링크: https://arxiv.org/pdf/2405.02803

단일 순방향 패스에서 Flash Attention의 수치 편차는 BF16의 Baseline Attention보다 훨씬 더 큰 것으로 나타났습니다.

구체적으로 이 방법은 다음을 포함하는 두 단계로 구성됩니다.

• 특정 최적화에서 수치 정확도를 교란시키는 마이크로 벤치마크 개발
• 데이터 기반 분석을 통해 수치 값 평가 Wasserstein 거리 기반 편향이 모델 가중치의 변화로 변환되는 방식입니다.

연구원들은 SOTA 최적화 기술인 Flash Attention을 분석하고 도입될 수 있는 수치 편차를 정량화했습니다. Flash Attention은 주의 메커니즘을 가속화하는 데 널리 사용되는 기술이며 종종 Transformer 모델에서 시스템 병목 현상으로 간주됩니다. Flash Attention은 속도를 향상시키고 메모리 액세스를 줄이면서도 알고리즘 최적화에 의존하므로 알고리즘 최적화로 인해 수치 편차가 증가할 수 있습니다.

연구원들은 재조정 요인을 추가하면 의도하지 않은 근사치가 발생하여 수치적 상충 관계가 발생하여 결과적으로 훈련 안정성에 영향을 미칠 수 있다는 가설을 세웠습니다.

다중 모드 텍스트-이미지 워크로드의 맥락에서 Flash Attention을 분석하여 Flash Attention과 기준선 간의 수치 편차의 잠재적 중요성을 확인했습니다. 궁극적으로 그들은 훈련 최적화의 수치적 편향과 그 후속 효과를 정량화하는 프레임워크를 도입합니다.

연구원들은 수치 편차를 정량화하는 데 다음과 같은 두 가지 주요 기여를 했습니다.

(1) 수치 정확도가 수치 편차에 미치는 영향을 분리하기 위해 마이크로 벤치마크를 설계했습니다.

연구원이 설계한 마이크로 벤치마크는 기존 블랙박스 최적화(예: Flash Attention)로 인해 발생하는 수치 편차를 측정하고 정량화하는 데 사용되는 기술입니다. 제공된 커널에서 일반적으로 사용할 수 없는 측면을 교란함으로써 낮은 수치 정밀도(BF16)에서 Flash Attention이 Baseline Attention에 비해 수치 편향이 대략 10배 더 높다는 발견을 개척했습니다.

(2) Wasserstein Distance 측정법을 기반으로 데이터 기반 분석을 수행했습니다.

이 분석을 통해 연구자들은 관찰된 수치 편차를 맥락화하고 이것이 다운스트림 모델 속성에 미치는 영향의 상한선을 형성합니다. 연구원의 사례 연구에서 그들은 관찰된 수치적 편향의 영향을 제한할 수 있었고 다음과 같은 사실을 발견했습니다. "Flash Attention은 저정밀 훈련의 약 1/2~1/5배인 모델 가중치 편향을 도입했습니다."

이 연구는 "수치 편향에 대한 교육 최적화의 영향을 정량화할 뿐만 아니라 맥락화"하기 위한 원칙적 접근 방식을 개발하는 것이 중요하다는 점을 강조합니다. , 훈련 불안정성) 측정하기 어려운 경우가 많습니다.

실험 방법

연구원들은 먼저 Flash Attention으로 인한 수치 편차를 분리하고 연구하기 위해 마이크로 벤치마크를 개발했습니다. 그림 2에서 볼 수 있듯이 그들은 다양한 수치 정밀도를 분석하고 알고리즘의 각 단계에서 잠재적인 최적화 조치를 적용하기 위해 Flash Attention을 수치적으로 다시 구현했습니다.

그림 2: Microbenchmark 설계 요약.

Flash Attention 코어는 현재 FP16 및 BF16 숫자 형식만 지원하기 때문에 이는 필요합니다. 이 커널은 CUDA 코드에 대한 래퍼 API 호출이기도 하므로 수치 편향의 영향을 조사하기 위해 알고리즘을 교란시키기가 어렵습니다.

반대로 마이크로 벤치마크 설계를 통해 알고리즘 내에서 정밀한 입력과 수정이 가능합니다. 연구원들은 원래 Flash Attention 커널에 대해 마이크로벤치마크를 검증했습니다.

그들은 모델 실행 중 각 단계에서 Attention 매트릭스의 출력을 비교하는 기술을 추가로 설계했습니다. 그리고 Attention이 호출될 때마다 Baseline Attention과 Flash Attention을 계산하도록 모델 코드를 수정하여 동일한 입력 행렬에 대한 정확한 출력 행렬 비교가 가능합니다.

이를 맥락에 맞게 설명하기 위해 Max Difference 및 Wasserstein Distance 측정항목을 사용하여 동일하고 독립적인 교육 실행을 통해 교육 전반에 걸쳐 모델 가중치의 차이를 정량화했습니다.

훈련 실험에서 연구원들은 텍스트 입력을 이미지로 변환하는 생성적 AI 워크로드(예: 텍스트-이미지 모델)를 사용했습니다. 그들은 Shutterstock 데이터 세트를 사용하여 모델을 재교육하고 NVIDIA 80GB A100 GPU 클러스터에서 실험을 실행했습니다.

마이크로 벤치마크를 통한 수치 편차 정량화

연구원들은 먼저 순방향 전달 프로세스에서 Flash Attention의 영향을 분석했습니다. 그들은 무작위로 초기화된 쿼리, 키 및 값 벡터가 동일하다는 조건 하에서 Attention으로 계산된 출력 행렬에 대한 다양한 수치 정밀도의 영향을 조사하기 위해 마이크로벤치마크를 사용했습니다.

그림 3에서 볼 수 있듯이 연구자가 BF16에서 FP64까지 다양한 숫자 형식을 사용할 때 Flash Attention과 Baseline Attention 사이의 수치 편차는 가수 자릿수가 증가함에 따라 감소합니다. 이는 수치적 차이가 더 적은 가수 자릿수를 갖는 데 내재된 근사치에 기인함을 시사합니다.

그림 3: 숫자 형식이 Flash Attention의 수치 편차에 미치는 영향.

이후 연구원은 표준 비교를 위해 FP64 수치 형식에서 Baseline Attention에 대한 "골든 값"을 설정한 다음 이 값과 다양한 수치 형식의 Attention 출력을 비교했습니다(그림 4 참조).

그림 4: FP64에서 기준 주의 "금 가치" 비교.

결과에 따르면 Flash Attention의 수치 편차는 BF16에서 Baseline의 수치 편차의 약 10배입니다.

이 관찰된 수치 편차를 추가로 분석하기 위해 연구원들은 타일 크기와 SRAM 크기를 일정하게 유지하면서 행렬의 시퀀스 길이를 스캔했습니다(그림 5 참조).

그림 5: Flash Attention의 수치 편차에 대한 시퀀스 길이의 영향.

그림에 표시된 것처럼 시퀀스 길이가 증가함에 따라 (a) 최대 차이의 상한으로 측정하거나 (b) 차이의 평균 및 표준 편차로 측정하면 Flash Attention과 Baseline의 차이가 나타납니다. 주의 수치 편차가 증가하고 있습니다.

또한 연구원들은 수치 편차의 영향을 더 잘 이해하기 위해 마이크로 벤치마크 설계를 사용하여 다양한 최적화로 실험을 수행합니다(그림 6 참조).

그림 6a는 블록 크기의 순서를 바꾸면 Flash Attention과 Baseline Attention 간의 수치 차이가 어떻게 증가하는지 보여줍니다. 타일 ​​크기를 정사각형으로 제한하는 등 그림 6b의 다른 교란은 수치 편향에 영향을 미치지 않습니다. 그림 6c는 블록/타일 크기가 클수록 수치 편차가 작아짐을 보여줍니다.

그림 6: 알고리즘 변경 및 관찰된 수치 편차에 미치는 영향.

가중치 차이를 통한 수치적 편향 이해

Flash Attention은 순방향 전달 동안 Attention 출력에 수치적 편향을 일으킬 수 있지만, 이 연구의 궁극적인 목표는 모델 훈련 중에 이것이 발생하는지 여부를 조사하는 데 어떤 효과가 있는지 확인하는 것입니다. 이는 훈련 불안정에 기여합니다.

따라서 연구원들은 Flash Attention이 훈련 중에 모델을 변경하는지, 즉 위에서 관찰한 Attention 출력의 차이가 훈련 중에 업데이트된 모델 가중치에 반영되는지를 정량화하기를 희망합니다.

연구원들은 두 가지 지표를 사용하여 Baseline Attention을 사용하여 훈련된 모델과 Flash Attention을 사용하여 훈련된 모델 간의 모델 가중치 차이를 측정했습니다. 최대 차이가 먼저 계산됩니다. 즉, 가중치 행렬 간 차이의 절대값을 구하고 최대값을 취하여 다음과 같이 편차의 상한을 얻습니다.

최대 차이는 다음과 같습니다. 수치 편차의 상한선이지만 각 행렬의 분포를 고려하지 않습니다. 따라서 연구자들은 텐서 간의 유사성을 측정하는 일반적인 척도인 Wasserstein Distance를 통해 가중치 차이를 정량화합니다. 계산적으로는 약간 더 복잡하지만 Wasserstein Distance에는 유사성을 측정하기 위한 텐서 분포의 모양 정보가 포함됩니다. 계산 공식은 다음과 같이 요약됩니다.

값이 낮을수록 행렬 간의 유사성이 높아집니다.

이 두 측정항목을 사용하여 연구원들은 훈련 과정 전반에 걸쳐 Flash Attention의 모델 가중치가 Baseline Attention과 비교하여 어떻게 변경되었는지 정량화했습니다.

Wasserstein 거리 및 최대 차이에 따라 전체 훈련 과정에서 Flash Attention을 추가하면 모델 가중치가 변경되고, 훈련이 계속됨에 따라 이 차이는 점점 더 커질 것입니다. 이는 Flash Attention을 사용하여 훈련한 모델이 Baseline Attention을 사용하여 훈련한 모델과 다르다는 것을 보여줍니다. 훈련된 동일한 모델이 다른 모델로 수렴되었습니다.

훈련은 확률론적 과정이며, 모델 구조의 특정 변화는 다운스트림 효과 및 정확성 측면에서 유사한 결과를 낳을 수 있습니다. 이는 Flash Attention과 Baseline Attention으로 훈련된 모델의 가중치가 다른 경우에도 주목할 만합니다.

모델을 완전히 훈련하고 정확성을 평가하는 것은 비용이 많이 들고 리소스 집약적인 작업입니다. 특히 훈련하는 데 수개월이 걸리는 대규모 모델의 경우 더욱 그렇습니다.

연구원은 탐색할 프록시를 구성했습니다.

(a) 이러한 체중 변화의 중요성은 무엇입니까?

(b) 이것이 널리 채택된 다른 훈련 최적화의 표준 체중 변경과 관련이 있을 수 있습니까?

이 목표를 달성하기 위해 연구원들은 다양한 시나리오에서 훈련 과정 동안 체중 차이가 어떻게 변하는지 비교하는 일련의 실험을 설계했습니다.

Flash Attention과 Baseline Attention을 사용한 훈련 과정을 비교한 것 외에도, 훈련 시작 시 가중치를 서로 다른 임의의 값으로 초기화한 동일한 훈련 과정에서 가중치의 차이도 수량화했습니다. 무작위 가중치 초기화는 일반적인 기술이고 종종 동일한 결과를 생성하므로 이는 한계를 제공합니다.

또한 연구원들은 다양한 정확도로 훈련된 모델 가중치의 변화도 측정했습니다. 수치 정밀도(예: FP16 대 FP32)는 Flash Attention 가중치의 중요성에 대한 상한선 역할을 하는 다운스트림 변경을 일으킬 가능성이 있습니다.

그림 8에서 볼 수 있듯이 Flash Attention을 사용한 모델 가중치 편향 변화율은 다양한 모델 초기화의 가중치 편향 변화율과 비슷하거나 작다는 것을 알 수 있습니다(빨간색과 파란색 곡선의 기울기에 유의하세요). .

또한, 다른 모델을 초기화할 때보다 FP16과 FP32를 사용할 때 무게 변화율이 더 높고 변화도 더 큽니다.

이러한 결과는 프록시를 제공하고 다음을 보여줍니다. "Flash Attention은 수치적 편향을 나타내지만 무작위 모델 초기화 및 낮은 정밀도 교육으로 인해 제한됩니다. 그리고 소개된 모델 가중치 편향은 낮은 정밀도로 교육할 때 약 10%입니다. . 1/2~1/5배. "

그림 8: Wasserstein 거리 측정법을 사용하여 측정한 상대적 체중 차이.

자세한 연구 내용은 원문을 참고해주세요.

위 내용은 Flash Attention은 안정적인가요? Meta와 Harvard는 모델 중량 편차가 ​​수십 배로 변동한다는 사실을 발견했습니다.의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!