저정밀도 학습 방법

새로운 유형의 하드웨어가 출시되면서 이를 보다 효율적으로 활용할 수 있는 새로운 학습 패러다임이 등장했습니다. 현재 이는 TransformersEngine(TE) 또는 torchao(네이티브 PyTorch FP8)와 같은 패키지를 사용하여 8bit 정밀도로 학습하는 형태로 나타납니다.

오늘 논의된 주제에 대한 소개로, 이 문서에서 정기적으로 참조할 예정이므로 low-precision usage guide를 검토해 보시는 것을 권장합니다.

간단한 차트

아래 차트는 학습 중 각 솔루션의 다른 비트 정밀도를 보여줍니다:

TransformersEngine

TransformersEngine는 8비트 부동소수점으로 학습을 시도하는 최초의 솔루션입니다. 모델의 특정 레이어를 FP8 엔진을 활용하여 비트 수(예: 32비트에서 8비트로)를 줄이면서도 모델의 최종 정확도를 저하시키지 않는 드롭인 대체 레이어로 교체하는 방식으로 작동합니다.

구체적으로, Accelerate는 다음 레이어를 찾아 TransformersEngine 버전으로 대체합니다:

• te.LayerNorm으로 nn.LayerNorm
• te.Linear으로 nn.Linear

그 결과 대부분의 레이어가 BF16로, 일부 레이어는 FP8로 구성되는 모델이 생성되어 메모리 사용량이 일부 감소합니다.

경험적으로, TransformerEngine를 사용할 때 성능 향상이 실제로 나타나기 시작하려면 모델의 레이어 대부분이 이 두 가지 레이어로 대체되어야 합니다. 그 결과, 파라미터 수가 수십억 개 수준이거나 그 이상인 대형 모델에서만 성능 개선이 확인되었습니다.

TransformerEngine는 FP8 계산 수행 방식과 그 동작을 사용자 정의할 수 있는 다양한 인수를 받을 수 있습니다. 인수의 전체 목록은 다음과 같습니다:

• margin: 그라디언트 스케일링에 사용할 마진입니다.
• interval: 스케일링 팩터를 다시 계산하는 주기에 사용할 간격.
• fp8_format``: The format to use for the FP8 recipe. Must be one ofHYBRIDorE4M3. (GenerallyHYBRIDfor training,E4M3` for evaluation)
• amax_history_len: 스케일링 계수 계산에 사용할 히스토리의 길이
• amax_compute_algo: 스케일링 계수 계산에 사용할 알고리즘. max 또는 most_recent 중 하나여야 합니다.
• override_linear_precision: fprop, dgrad, wgrad GEMM을 더 높은 정밀도로 실행할지 여부.

모델의 성능을 최적화하는 데 도움이 되도록 utils.TERecipeKwargs의 일부로 이 항목들을 각각 사용자 정의할 수 있습니다.

앞서 언급한 차트를 살펴보면, TE는 계산 레이어만 FP8로 캐스팅하고 나머지는 모두 FP32로 유지합니다. 이로 인해 가장 많은 메모리를 사용하지만, 학습 중 최종 정확도 손실이 최소로 보장된다는 장점이 있습니다.