Hugging Face GPU 추론 최적화

Hugging Face GPU 추론 최적화 (Performance on GPU)

GPU는 메모리 대역폭과 병렬성에 최적화되어 있어서 머신러닝의 표준 하드웨어예요. 모델이 점점 커지면서 GPU가 효율적으로 최고 성능을 내도록 만드는 것이 더 중요해졌어요. 이 가이드는 GPU에서 추론을 최적화하는 몇 가지 방법을 보여줘요. 아래 최적화 방법들은 서로 결합할 수 있고 분산 GPU에서도 동작해요.

출처: Hugging Face GPU 성능 문서

bitsandbytes (양자화)

bitsandbytes는 8-bit와 4-bit 양자화를 지원하는 라이브러리예요. 양자화는 원래 전체 정밀도보다 낮은 정밀도로 가중치를 표현해서 메모리 요구량을 줄이고 큰 모델을 메모리에 맞게 해줘요.

pip install bitsandbytes accelerate

8-bit 양자화로 텍스트 생성할 때는 고수준 Pipeline API 대신 generate()를 써야 해요. Pipeline은 8-bit 모델에 최적화되어 있지 않아 느리고, 일부 샘플링 전략(nucleus sampling)도 지원하지 않아요.

BitsAndBytesConfig로 load_in_8bit=True를 설정하고, from_pretrained()의 quantization_config에 넘겨요. device_map="auto"로 Accelerate가 모델을 하드웨어에 자동 분산하게 하고 입력을 모델과 같은 장치에 두세요.

from transformers import BitsAndBytesConfig, AutoTokenizer, AutoModelForCausalLM

quantization_config = BitsAndBytesConfig(load_in_8bit=True)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B",
                                             device_map="auto",
                                             quantization_config=quantization_config)

분산 설정에서는 max_memory 파라미터로 각 GPU에 할당할 메모리를 매핑해요.

max_memory_mapping = {0: "16GB", 1: "16GB"}
model_8bit = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    device_map="auto",
    quantization_config=quantization_config,
    max_memory=max_memory_mapping,
)

Optimum (ONNX Runtime)

Optimum은 다양한 하드웨어에서 모델 성능을 최적화하는 Hugging Face 라이브러리예요. ONNX Runtime(ORT)을 지원하며, 공통 연산을 단일 노드로 퓨전하고 상수 폴딩으로 계산 수를 줄여요. 가장 계산 집약적인 연산은 GPU에, 나머지는 CPU에 배치해 지능적으로 분산해요.

from optimum.onnxruntime import ORTModelForSequenceClassification

ort_model = ORTModelForSequenceClassification.from_pretrained(
    "distilbert/distilbert-base-uncased-finetuned-sst-2-english",
    provider="CUDAExecutionProvider",
)

provider는 하드웨어에 맞게 설정해요: CUDAExecutionProvider(CUDA GPU), ROCMExecutionProvider(AMD), TensorrtExecutionProvider(TensorRT).

Attention 인터페이스 (SDPA)

PyTorch의 torch.nn.functional.scaled_dot_product_attention(SDPA)는 scaled dot product attention의 네이티브 구현으로, 트랜스포머의 attention보다 효율적이에요. v2.1.1 이상 PyTorch에서 구현이 가능하면 기본으로 사용돼요. 명시적으로 attn_implementation="sdpa"로 설정하거나 model.set_attention_implementation("sdpa")로 활성화할 수 있어요.

FlashAttention2

FlashAttention을 설치하고 attn_implementation="flash_attention_2"로 설정하면 긴 시퀀스에서 추론 속도를 크게 높일 수 있어요.

pip install flash-attn --no-build-isolation

FlashAttention2는 fp16 또는 bf16 타입만 지원하므로 먼저 타입을 캐스팅해야 해요.

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    device_map="auto",
    dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)

FlashAttention2는 패딩 토큰과 attention score를 계산하는 것을 지원하지 않아서, 배치 추론에 패딩 토큰이 있으면 수동으로 패딩/언패딩해야 해요. 그래서 짧은 시퀀스에서는 오히려 오버헤드로 느려질 수 있어요.

더 알아보기