Hugging Face Transformers, TensorFlow 및 TPU를 사용한 언어 모델 훈련

Hugging Face는 TensorFlow와 TPU 팟을 사용해 처음부터 마스크 언어 모델을 훈련하기 위한 확장 가능한 종단 간 가이드를 제공하며, 고성능 컴퓨팅을 위해 XLA 호환성을 활용합니다.

XLA 호환성 및 TPU 접근성

TPU에서 TensorFlow 모델을 훈련하는 것은 과거에 XLA 호환성 부족과 비네이티브 TensorFlow 연산에 의존하는 데이터 콜레이터 때문에 어려웠습니다. Hugging Face는 코드베이스를 업데이트하여 대부분의 TensorFlow 모델이 XLA와 호환되도록 하여 이러한 장애물을 제거하고 TPU 훈련을 보다 쉽게 접근할 수 있게 만들었습니다.

고성능 GPU가 지속적으로 부족한 상황에서 이 변화는 매우 중요합니다. TPU는 초고성능 컴퓨팅 하드웨어에 접근할 수 있는 고성능 대안을 제공하여, GPU 가용성에만 의존하지 않고 대규모 생성 AI 및 LLM을 훈련할 수 있는 확장 가능한 경로를 제공합니다.

종단 간 훈련 워크플로우

이 접근법의 확장성을 보여주기 위해, Hugging Face는 WikiText (v1) 데이터셋을 사용해 RoBERTa-base 모델을 처음부터 훈련했습니다. 이 과정은 TPU 하드웨어에서 효율성을 보장하기 위해 특정 파이프라인을 따릅니다:

1. 토크나이저 훈련 및 데이터 준비

모델을 처음부터 훈련하기 때문에 맞춤형 토크나이저가 필요합니다. 워크플로우는 다음을 포함합니다:

  • WikiText 데이터셋의 train 분할을 ‚datasets를 통해 로드합니다.
  • ‚tokenizers를 사용해 Unigram 모델을 훈련합니다.
  • 생성된 토크나이저를 Hugging Face Hub에 업로드합니다.

2. TFRecord 샤드 생성

대규모 병렬 처리를 가능하게 하기 위해 데이터를 단일 파일이 아닌 TFRecord 샤드로 변환합니다. 토크나이징 전략은 샘플을 연결하고 고정 크기 청크(128 토큰)로 나누어 잘림으로 인한 텍스트 손실을 방지하는 것입니다.

이 샤드들은 Google Cloud Storage(GCS) 버킷에 업로드됩니다. 이는 호스트 메모리가 제한된 TPU 노드가 데이터를 GCS에서 직접 스트리밍해야 하기 때문에 필요합니다. (참고: TPU VM은 로컬 데이터셋이나 영구 스토리지를 사용할 수 있습니다.)

3. 모델 초기화 및 분산 훈련

데이터 병렬 처리를 사용해 TPU 워커에 훈련을 분산하려면 모델과 옵티마이저를 TPUStrategy 스코프 내에서 초기화해야 합니다:

import tensorflow as tf

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(...)
strategy = tf.distribute.TPUStrategy(tpu)

with strategy.scope():
    # Model and tokenizer initialization happens here
    model = TFAutoModelForMaskedLM.from_config(config)

TPU 훈련을 위한 핵심 기술 요구사항

TPU를 성공적으로 통합하려면 데이터 파이프라인과 모델 설정에 특정 구성이 필요합니다:

  • TensorFlow-네이티브 데이터 콜레이터: DataCollatorForLanguageModeling은 return_tensor="tf"로 설정해야 합니다. 이는 콜레이터가 NumPy 배열이 아닌 TensorFlow 텐서를 반환하도록 하여 TPU 호환성에 필수적입니다.
  • GCS 통합: TensorFlow의 tf.io.gfile.glob은 gs:// 식별자를 사용해 GCS 버킷에서 TFRecord 샤드를 원활하게 읽을 수 있게 합니다.
  • 모델 체크포인트: PushToHubCallback은 훈련 중 모델 체크포인트를 Hugging Face Hub에 직접 동기화하는 데 사용됩니다.

추론 및 결과

훈련이 완료되면, 표준 Hugging Face pipeline API에 framework="tf" 인자를 사용해 모델을 추론용으로 배포할 수 있습니다. 훈련된 RoBERTa-base 모델은 1e-4 학습률로 더 긴 기간 동안 훈련되었으며, 결과 가중치는 Hugging Face Hub에서 확인할 수 있습니다.

Sources