※본포스팅은 Nvidia Transformer Engine Docs를 참고했습니다.
Transformer Engine으로 FP8 사용하기
H100에서는 FP8(8-bit floating point) 형식을 지원한다
Introduction to FP8
H100에서 두 가지 FP8 형식을 지원한다.
1) E4M3 - +/- 448
2) E5M2 - +/- 57344 Dynamic Range가 넓지만 정확도가 떨어짐.
<Training pass별 FP8 Type 선>
Forward pass - E4M3 방식 사용 , weight 값을 계산 → 정확도 중요
Backward pass - E5M2 방식 사용, Gradient 값 계산 → 넓은 Dynamic Range 필요
Mixed Precision Training
FP16 작동방법을 통해서 FP8 작동 방식 이해해보자.
FP8을 적용하기 위해서는 FP16 학습 방식에서 어떤점을 유의해서 FP32사용할때와의 격차를 극복했는지
알아볼 필요가 있다. FP16이 선배격인셈. 앞선자들의 자취를 따라가는게 가장 빨리 더 나은 방법을 생각해내는 방법이다.
1) FP16 사용시에도 정확도가 떨어질 수 있다. 이때는 행렬연산, convolution 연산등 실행하는 연산에 따라
output의 값의 범위를 체크해봐야 한다.
2) Gradient 계산시에는 OverFlow, UnderFlow 발생에 유의해야 한다. 이를 방지하기위해서는 Dynamic Loss Scaling 을
통해 Loss값을 정규분포로 만들어 주는 방식이 효과적이다.
Mixed Precision Training with FP8
FP8의 경우를 보자. FP8은 Activation연산, Gradient 계산시에는 충분하지만, 다른 모든 경우에 적합하지는 않는다.
Loss scaling도 한가지 방식으로 모두 적용하는게 아닌 FP8 텐서마다 다르게 적용하는게 필요하다.
*infeasible 실행불가능한
*distinct 분명한, 뚜렷한
Scaling factor 선택 전략
- 여기서 Factor는 0~1사이의 값으로, 그래디언트값이 커지지 않도록 곱해주는 상수를 의미
- Just-in-time
- 이 방법은 출력텐서(output)의 절대값의 최대값(amax)을 기준으로 스케일링 요소를 선택한다. 매번 그래디언트의 최대값을 측정하고 그 값으로 스케일링 Factor를 결정한다. 이렇게 되면 매번 최대값을 측정해야 하므로 많은 오버헤드가 발생하며 FP8 사용의 이점이 줄어든다.
- Delayed Scaling
- 그래디언트의 스케일링 요소를 이전 일정한 iteration들을 고려하여 계산하는 방법입니다. 이 방법을 사용하면 FP8 연산의 최대 성능을 발휘할 수 있지만, 그래디언트의 최댓값들의 기록을 저장하기 위해 FP8 연산자들의 추가적인 매개변수로 기록을 유지해야 합니다.
- 요약
- just-in-time 방식은 매번 그래디언트를 연산할때마다 최대값을 측정하는거고,
- Delayed Scaling은 이제까지 연산한것 중에 최대값만 관리해서 Factor를 사용함
- 결론은 Delayed Scaling가 메모리를 덜잡아먹는다. 매번 최대값을 측정 안해줘도 되니까
Using FP8 with Transformer Engine
- Transformer Engine은 FP8을 쉽게 사용할 수 있도록 라이브러리로 제공한다.
FP8 사용방안(recipe)
라이브러리 패키지 transformer_engine.common.recipe 의 DelayedScaling 모듈은 FP8 학습에서 필요한 모든
옵션을 저장한다. 이걸로 Scaling Factor에 필요한 절대값을 계산한다.
FP8 학습의 config 작성은 다음과 같다.
from transformer_engine.common.recipe import Format, DelayedScaling
fp8_format = Format.HYBRID # E4M3 during forward pass, E5M2 during backward pass
fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo="max")FP8 자동변환(Autocasting)
**Cast
- (금속, 플라스틱 등을) 주형에 부어 만들다: 주조(鑄造)의 의미로 쓰이며, 금속이나 플라스틱 등을 유리한 형에 녹여 부어서 만드는 공정을 가리킵니다.
- (연기, 빛 등을) 내뿜다: 어떤 물체로부터 연기, 빛 등이 퍼져 나가는 동작을 말합니다.
- (역할, 역활을) 맡다: 연기자나 배우가 연극, 영화 등에서 특정한 역할을 맡는 것을 의미합니다.
- (눈총, 던져서) 던지다: 무엇을 손이나 무기로 던지거나 투척하는 동작을 가리킵니다.
- (시선, 눈 등을) 돌리다: 시선, 눈 등을 다른 방향으로 돌려 보거나 이동시키는 것을 의미합니다.
- (무엇을) 결정하다, 정하다: 어떤 일에 대해 결정을 내리거나 정하는 것을 의미합니다.
- (마법, 주문 등으로) 변하게 하다: 마법이나 주문을 사용하여 물건을 변하게 하거나 마음을 조종하는 것을 의미합니다.
- FP8이 모든 연산에 적용되는 것은 아니다.
- TE 라이브러리 모든 모듈은 FP8을 최대한 활용해서 사용되면서도 정확도는 유지되도록 설계되었다.
- FP8연산을 사용하려면 fp8_autocast 에 감싸는 방식으로 코드를 작성해 줘야 한다.
import transformer_engine.pytorch as te
import torch
torch.manual_seed(12345)
my_linear = te.Linear(768, 768, bias=True)
inp = torch.rand((1024, 768)).cuda()
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
out_fp8 = my_linear(inp)이러한 단순한 방식은 아래와 같은 과정을 자동으로 처리해준다.
- 모든 Input을 FP8로 변환한다.
- 최대값을 업데이트 한다.
- Scaling Factor를 계산한다.
주의!
Transformer Engine의 Linear 레이어에서 FP8을 지원하는 것은 현재 두 차원 모두가 16로 나누어 떨어지는 형태의 텐서에 한정됩니다. 전체 Transformer 네트워크의 입력에 대해 이를 적용하기 위해서는 일반적으로 시퀀스 길이를 16의 배수로 패딩해야 합니다.
Backward Pass 처리
- fp8_autocast 영역 내에서 모델을 실행할 때, 특히 다중 GPU 학습에서는 스케일링 요소와 Amax history를 동기화하기 위해 일부 통신이 필요합니다. 이러한 통신을 오버헤드를 최소화하면서 수행하기 위해 fp8_autocast 컨텍스트 매니저는 통신 전에 텐서들을 집계합니다.
- 이러한 집계로 인해 역전파 호출은 fp8_autocast 컨텍스트 매니저 밖에서 발생해야 합니다. 이는 계산 정확성에 영향을 미치지 않습니다. 역전파의 정밀도는 순전파의 정밀도에 의해 결정됩니다.
loss_fp8 = out_fp8.mean()
loss_fp8.backward() # This backward pass uses FP8, since out_fp8 was calculated inside fp8_autocast
out_fp32 = my_linear(inp)
loss_fp32 = out_fp32.mean()
loss_fp32.backward() # This backward pass does not use FP8, since out_fp32 was calculated outside fp8_autocast
정확도(Precision)
FP8과 FP32를 비교하면 값이 근사하지만 다르다는 걸 확인 할 수 있다.
out_fp8
tensor([[ 0.2276, 0.2627, 0.3001, ..., 0.0346, 0.2211, 0.1188],
[-0.0963, -0.3725, 0.1717, ..., 0.0901, 0.0522, -0.3472],
[ 0.4526, 0.3482, 0.5976, ..., -0.0687, -0.0382, 0.1566],
...,
[ 0.1698, 0.6061, 0.0385, ..., -0.2875, -0.1152, -0.0260],
[ 0.0679, 0.2946, 0.2751, ..., -0.2284, 0.0517, -0.1441],
[ 0.1865, 0.2353, 0.9172, ..., 0.1085, 0.1135, 0.1438]],
device='cuda:0', grad_fn=<_LinearBackward>)
out_fp32
tensor([[ 0.2373, 0.2674, 0.2980, ..., 0.0233, 0.2498, 0.1131],
[-0.0767, -0.3778, 0.1862, ..., 0.0858, 0.0676, -0.3369],
[ 0.4615, 0.3593, 0.5813, ..., -0.0779, -0.0349, 0.1422],
...,
[ 0.1914, 0.6038, 0.0382, ..., -0.2847, -0.0991, -0.0423],
[ 0.0864, 0.2895, 0.2719, ..., -0.2388, 0.0772, -0.1541],
[ 0.2019, 0.2275, 0.9027, ..., 0.1022, 0.1300, 0.1444]],
device='cuda:0', grad_fn=<_LinearBackward>)
이런 현상이 발생하는 이유는 FP8 경우 입력과 가중치가 연산 전에 모두 FP8로 변환되기 때문입니다. 원래의 입력 대신 FP8로 표현 가능한 입력을 사용한다면 이를 확인할 수 있습니다 (quickstart_utils.py에 정의된 함수를 사용하여).
from quickstart_utils import cast_to_representable
inp_representable = cast_to_representable(inp)
my_linear.weight.data = cast_to_representable(my_linear.weight.data)
out_fp32_representable = my_linear(inp_representable)
print(out_fp32_representable)tensor([[ 0.2276, 0.2629, 0.3000, ..., 0.0346, 0.2211, 0.1188],
[-0.0963, -0.3724, 0.1717, ..., 0.0901, 0.0522, -0.3470],
[ 0.4526, 0.3479, 0.5976, ..., -0.0686, -0.0382, 0.1566],
...,
[ 0.1698, 0.6062, 0.0385, ..., -0.2876, -0.1152, -0.0260],
[ 0.0679, 0.2947, 0.2750, ..., -0.2284, 0.0516, -0.1441],
[ 0.1865, 0.2353, 0.9170, ..., 0.1085, 0.1135, 0.1438]],
device='cuda:0', grad_fn=<_LinearBackward>)
두개의 값의 차를 계산해보면 매우 작다.
out_fp8 - out_fp32_representable
tensor([[ 4.9591e-05, -1.9073e-04, 9.5367e-05, ..., -3.8147e-06,
4.1962e-05, 2.2888e-05],
[ 2.2888e-05, -3.4332e-05, 2.2888e-05, ..., 2.6703e-05,
5.3406e-05, -1.4114e-04],
[-3.8147e-05, 2.6703e-04, -3.8147e-06, ..., -5.7220e-05,
4.1962e-05, -1.9073e-05],
...,
[ 1.1444e-05, -7.2479e-05, -3.8147e-06, ..., 5.3406e-05,
-1.5259e-05, 2.2888e-05],
[ 4.9591e-05, -9.5367e-05, 6.8665e-05, ..., -1.5259e-05,
7.6294e-05, 4.5776e-05],
[-1.5259e-05, -7.6294e-06, 1.8692e-04, ..., -3.0518e-05,
-4.5776e-05, 7.6294e-06]], device='cuda:0', grad_fn=<SubBackward0>)FP8 실행에서 오는 결과의 차이는(정확도) 훈련 과정에서 중요하지 않지만, 모델을 디버깅하는 동안 그것들을 이해하는 것이 좋다.