TensorFlow 모델을 위한 XLA 통합
XLA(Accelerated Linear Algebra)는 TensorFlow 모델의 실행 시간을 가속화하기 위한 컴파일러입니다. 공식 문서에 따르면 다음과 같습니다:
XLA(Accelerated Linear Algebra)는 선형 대수를 위한 도메인 특화 컴파일러로, TensorFlow 모델을 소스 코드 변경 없이 가속화할 수 있습니다.
TensorFlow에서 XLA를 사용하는 것은 간단합니다. XLA는 tensorflow 라이브러리 내에 패키지로 제공되며, tf.function과 같은 그래프 생성 함수에서 jit_compile 인수를 사용하여 활성화할 수 있습니다. fit() 및 predict()와 같은 Keras 메소드를 사용하는 경우, jit_compile 인수를 model.compile()에 전달하여 XLA를 간단하게 활성화할 수 있습니다. 그러나 XLA는 이러한 메소드에 국한되지 않고 임의의 tf.function을 가속화하는 데에도 사용할 수 있습니다.
🤗 Transformers에서는 GPT2, T5, OPT와 같은 모델의 텍스트 생성, 그리고 Whisper와 같은 모델의 음성 처리를 포함하여 여러 TensorFlow 메소드가 XLA와 호환되도록 다시 작성되었습니다.
정확한 속도 향상은 모델에 따라 다르지만, 🤗 Transformers 내의 TensorFlow 텍스트 생성 모델의 경우 최대 100배의 속도 향상을 확인했습니다. 이 문서에서는 이러한 모델에 대해 XLA를 사용하여 최대 성능을 얻는 방법을 설명합니다. 또한 XLA 통합의 벤치마크 및 디자인 철학에 대한 추가 자료 링크도 제공할 것입니다.
XLA를 사용하여 TF 함수 실행하기
TensorFlow에서 다음과 같은 모델을 고려해 봅시다:
import tensorflow as tf
model = tf.keras.Sequential(
[tf.keras.layers.Dense(10, input_shape=(10,), activation="relu"), tf.keras.layers.Dense(5, activation="softmax")]
)위 모델은 차원이 (10, )인 입력을 받습니다. 다음과 같이 모델을 사용하여 순전파를 실행할 수 있습니다:
# 모델에 대한 임의의 입력을 생성합니다.
batch_size = 16
input_vector_dim = 10
random_inputs = tf.random.normal((batch_size, input_vector_dim))
# 순전파를 실행합니다.
_ = model(random_inputs)XLA로 컴파일된 함수로 순전파를 실행하려면 다음과 같이 해야 합니다:
xla_fn = tf.function(model, jit_compile=True)
_ = xla_fn(random_inputs)model의 기본 call() 함수는 XLA 그래프를 컴파일하는 데 사용됩니다. 그러나 다른 모델 함수를 XLA로 컴파일하려면 다음과 같이 할 수도 있습니다:
my_xla_fn = tf.function(model.my_xla_fn, jit_compile=True)🤗 Transformers에서 XLA를 사용하여 TF 텍스트 생성 모델 실행하기
🤗 Transformers에서 XLA로 가속화된 생성을 활성화하려면 최신 버전의 transformers가 설치되어 있어야 합니다. 다음과 같이 설치할 수 있습니다:
pip install transformers --upgrade
그리고 다음 코드를 실행할 수 있습니다:
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM
# 최소 버전의 Transformers가 설치되어 있지 않다면 오류가 발생합니다.
from transformers.utils import check_min_version
check_min_version("4.21.0")
tokenizer = AutoTokenizer.from_pretrained("gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
input_string = ["TensorFlow is"]
# XLA 생성 함수를 만들기 위한 한 줄
xla_generate = tf.function(model.generate, jit_compile=True)
tokenized_input = tokenizer(input_string, return_tensors="tf")
generated_tokens = xla_generate(**tokenized_input, num_beams=2)
decoded_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
print(f"Generated -- {decoded_text}")
# Generated -- TensorFlow is an open-source, open-source, distributed-source application # framework for the알 수 있듯이, generate()에서 XLA를 활성화하는 것은 단 한 줄의 코드입니다. 코드의 나머지 부분은 변경되지 않습니다. 그러나 위 코드 스니펫에서는 XLA에 특정한 몇 가지 주의할 점이 있습니다. XLA가 가져다줄 속도 향상을 실현하기 위해서는 이를 알고 있어야 합니다. 다음 섹션에서 이에 대해 논의합니다.
주의할 점
XLA 활성화 함수(xla_generate()와 같은)를 처음 실행할 때 내부적으로 계산 그래프를 추론하려고 하며, 이는 시간이 소요됩니다. 이 과정은 “추적(tracing)”이라고 알려져 있습니다.
생성 시간이 빠르지 않다는 것을 알 수 있을 것입니다. xla_generate()(또는 다른 XLA 활성화 함수)의 연속 호출은 함수에 전달된 입력이 초기에 구축된 계산 그래프와 동일한 형태를 따른다면, 계산 그래프를 추론할 필요가 없습니다. 이는 입력 형태가 고정된 모달리티(예: 이미지)에는 문제가 되지 않지만, 가변 입력 형태 모달리티(예: 텍스트)를 사용할 때 주의해야 합니다.
xla_generate()가 항상 동일한 입력 형태로 동작하도록 하려면, 토크나이저를 호출할 때 padding 인수를 지정할 수 있습니다.
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
input_string = ["TensorFlow is"]
xla_generate = tf.function(model.generate, jit_compile=True)
# 여기서, padding 옵션이 있는 토크나이저를 호출합니다.
tokenized_input = tokenizer(input_string, pad_to_multiple_of=8, padding=True, return_tensors="tf")
generated_tokens = xla_generate(**tokenized_input, num_beams=2)
decoded_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
print(f"Generated -- {decoded_text}")이렇게 하면 xla_generate()에 대한 입력이 항상 추적된 형태로 전달되어 생성 시간이 가속화됩니다. 다음 코드로 이를 확인할 수 있습니다:
import time
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
xla_generate = tf.function(model.generate, jit_compile=True)
for input_string in ["TensorFlow is", "TensorFlow is a", "TFLite is a"]:
tokenized_input = tokenizer(input_string, pad_to_multiple_of=8, padding=True, return_tensors="tf")
start = time.time_ns()
generated_tokens = xla_generate(**tokenized_input, num_beams=2)
end = time.time_ns()
print(f"Execution time -- {(end - start) / 1e6:.1f} ms\n")Tesla T4 GPU에서는 다음과 같은 출력을 예상할 수 있습니다:
Execution time -- 30819.6 ms Execution time -- 79.0 ms Execution time -- 78.9 ms
xla_generate()의 첫 번째 호출은 추적 때문에 시간이 오래 걸리지만, 연속 호출은 몇 배나 빠릅니다. 생성 옵션에 대한 어떤 변경이든 다시 추적을 유발하므로 생성 시간이 느려질 수 있음을 명심하세요.
이 문서에서는 🤗 Transformers에서 제공하는 모든 텍스트 생성 옵션을 다루지 않았습니다. 고급 사용 사례에 대해 문서를 참조하시기 바랍니다.
추가 자료
여기에 🤗 Transformers와 XLA에 대해 더 자세히 알고 싶은 경우 도움이 될 수 있는 몇 가지 추가 자료를 제공합니다.
- 이 Colab 노트북은 XLA와 호환되는 인코더-디코더(T5와 같은) 및 디코더 전용(GPT2와 같은) 텍스트 생성 모델을 실험해 볼 수 있는 대화형 데모를 제공합니다.
- 이 블로그 글은 TensorFlow에서 XLA에 대한 친절한 소개와 함께 XLA와 호환되는 모델의 비교 벤치마크에 대한 개요를 제공합니다.
- 이 블로그 글은 🤗 Transformers의 TensorFlow 모델에 XLA 지원을 추가하는 것에 대한 디자인 철학을 논의합니다.
- XLA와 TensorFlow 그래프에 대해 더 자세히 알고 싶은 경우 추천하는 글: