TensorFlow와 XLA를 이용한 빠른 텍스트 생성

TL;DR

Hugging Face는 TensorFlow용 transformers 라이브러리에서 텍스트 생성을 위한 XLA(Accelerated Linear Algebra) 컴파일을 지원하도록 만들었습니다. 이 최적화는 생성 속도를 최대 100배까지 끌어올릴 수 있으며, 많은 벤치마크에서 텍스트 생성 작업에 있어 PyTorch보다 뛰어난 성능을 보입니다.

XLA로 TensorFlow 가속하기

XLA는 TensorFlow 모델을 가속화하도록 설계된 컴파일러이며, JAX와 일부 PyTorch 구현의 기반이 되기도 합니다. TensorFlow 2는 투명성과 디버깅 편의를 위해 Eager Execution을 사용하지만, 이로 인해 그래프 모드에서 얻을 수 있던 일부 성능 이점을 잃게 됩니다. 이러한 이점을 회복하려면 사용자가 tf.function으로 함수를 래핑하면 코드가 그래프로 변환됩니다.

tf.function이나 tf.keras.Model.compile에 jit_compile=True 인자를 추가하면 XLA 컴파일을 트리거할 수 있습니다. XLA로 컴파일된 함수는 첫 호출 시 컴파일 과정 때문에 느리지만, 동일한 텐서 형태와 타입을 가진 이후 호출은 크게 빨라집니다.

XLA 텍스트 생성 구현 요구사항

XLA는 JIT(Just‑In‑Time) 컴파일과 다형성을 기반으로 합니다. 텍스트 생성 중에 비용이 많이 드는 재컴파일(트레이싱)을 방지하려면 다음 기술적 요구사항을 충족해야 합니다:

입력 패딩

XLA는 텐서 형태, 타입, 혹은 텐서가 아닌 인자가 달라질 때마다 새로운 컴파일 단계를 시작합니다. 따라서 입력 프롬프트는 일관된 길이로 패딩해야 합니다. Hugging Face는 pad_to_multiple_of 인자를 토크나이저 클래스에 사용해 입력 유연성을 유지하면서 가능한 형태 수를 제한할 것을 권장합니다.

코드베이스 벡터화

자동 회귀 텍스트 생성은 본질적으로 동적이며, 텐서를 확장하거나 동적 슬라이스를 사용하는 경우가 많아 XLA에 부적합합니다. XLA 지원을 가능하게 하기 위해 Hugging Face는 TensorFlow 텍스트 생성 코드베이스를 벡터화하고 패딩된 고정 크기 구조를 활용하도록 재작성했습니다. 또한, NLP 모델을 수정해 위치 임베딩이 이러한 패딩 구조에서도 올바르게 동작하도록 했습니다.

Transformers의 텍스트 생성 기능

transformers 라이브러리의 generate 함수는 여러 디코딩 전략을 지원합니다:

  • Greedy Decoding: 기본 결정적 접근 방식(do_sample=False)으로, 각 단계에서 가장 가능성이 높은 토큰을 선택합니다.
  • Sampling: 확률적 접근 방식(do_sample=True)이며, temperature 설정을 통해 무작위성을 제어합니다. 낮은 값은 높은 확률 토큰을 우선시하고, 높은 값은 엔트로피를 증가시킵니다.
  • Beam Search: num_beams가 1보다 클 때 활성화되며, 탐색을 통해 높은 확률 시퀀스를 찾아 Greedy Decoding보다 출력 품질을 향상시킵니다.

성능 벤치마크

여러 GPU 모델에 걸친 TensorFlow와 PyTorch 비교 벤치마크는 두 가지 주요 결과를 보여줍니다:

  1. 엄청난 속도 향상: XLA를 사용할 경우 TensorFlow 텍스트 생성 속도가 크게 빨라지며, 경우에 따라 100배 이상의 속도 향상을 기록합니다.
  2. 프레임워크 비교: 대부분의 경우 XLA가 적용된 TensorFlow가 가장 빠른 옵션이며, 텍스트 생성 작업에서 PyTorch보다 최대 9배까지 빠르게 동작합니다.

Sources