Transformers documentation

TensorFlow 모델을 위한 XLA 통합

You are viewing v4.44.2 version. A newer version v4.46.0 is available.
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

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("openai-community/gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("openai-community/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("openai-community/gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("openai-community/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("openai-community/gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("openai-community/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에 대해 더 자세히 알고 싶은 경우 도움이 될 수 있는 몇 가지 추가 자료를 제공합니다.

< > Update on GitHub