beautifulsoups.dev
Published on

Transformers 로컬 LLM 느림·OOM, 4bit+FlashAttn2

Authors

로컬 GPU에서 transformers로 LLM을 올리다 보면 두 가지 문제가 반복됩니다.

  • 로딩이 지나치게 느림: 모델 가중치 다운로드·디스크 I/O·샤딩 로딩·CPU 오프로딩 때문에 “한참 멈춘 것처럼” 보입니다.
  • OOM(Out Of Memory): 모델 본체는 겨우 올라가도, 프롬프트가 길어지거나 배치가 늘면 KV 캐시가 폭증해 바로 터집니다.

이 글은 4bit 양자화(bitsandbytes) + FlashAttention2 조합을 중심으로, 로딩 속도와 VRAM 사용량을 동시에 줄이는 방법을 실전 코드로 정리합니다.

관련해서 KV 캐시까지 포함해 OOM을 더 촘촘히 막고 싶다면 아래 글도 함께 참고하세요.

왜 로컬 로딩이 느려질까: 병목 체크리스트

체감 로딩 속도는 “GPU 연산”보다 아래 항목에서 더 많이 갈립니다.

1) 모델 파일 다운로드/검증과 디스크 I/O

  • 첫 실행은 huggingface_hub가 수 GB~수십 GB를 내려받고, 샤드 파일을 여러 개 열어 읽습니다.
  • NVMe가 아닌 SATA SSD/HDD에서는 샤드 로딩이 길어집니다.

대응

  • 캐시 디렉터리를 빠른 디스크로 옮기기: HF_HOME, TRANSFORMERS_CACHE
  • 사내/오프라인 환경이라면 미리 snapshot_download로 받아두기

2) device_map="auto"의 CPU 오프로딩

VRAM이 빡빡하면 accelerate가 일부 레이어를 CPU로 내리는데, 이 경우

  • 로딩도 느려지고
  • 추론도 PCIe 왕복 때문에 급격히 느려집니다.

대응

  • 4bit로 VRAM을 줄여 GPU에 더 많이 상주시킨다
  • 필요 시 max_memory로 의도적으로 맵을 조정한다

3) dtype/커널 미스매치

  • Ampere 이상이면 보통 bfloat16 또는 float16이 유리합니다.
  • FlashAttention2를 쓰려면 보통 CUDA/torch/flash-attn 버전 궁합이 맞아야 합니다.

OOM의 진짜 범인: 모델 가중치 vs KV 캐시

OOM을 “모델이 커서”라고만 생각하면 해결이 늦어집니다.

  • 모델 가중치 VRAM: 로딩 시점에 크게 한 번 먹습니다.
  • KV 캐시 VRAM: 추론 중 토큰이 늘어날수록 선형(혹은 배치까지 곱해져) 증가합니다.

즉, 4bit로 모델을 줄여도 **컨텍스트 길이(max_new_tokens, 프롬프트 길이, max_position_embeddings)**를 무리하게 쓰면 KV 캐시로 터질 수 있습니다.

핵심 해법 1: 4bit 양자화로 “일단 올려라”

bitsandbytes 4bit는 로컬 LLM에서 가장 즉효가 큽니다.

  • 장점: VRAM을 크게 줄여 CPU 오프로딩을 피하기 쉬움
  • 단점: 약간의 품질 저하 가능(대부분의 일반 추론에서는 체감이 작거나 허용 가능)

아래는 가장 무난한 조합인 NF4 + double quant + bfloat16 compute 예시입니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "meta-llama/Llama-2-7b-chat-hf"  # 예시

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

model.eval()

4bit에서 자주 하는 실수

  1. torch_dtype를 float32로 두기
  • 4bit여도 일부 연산 dtype가 커지면 VRAM/속도에 손해입니다.
  1. device_map 없이 통째로 GPU에 올리려다 OOM
  • VRAM이 애매하면 device_map="auto"로 시작하되, 목표는 “오프로딩 최소화”입니다.
  1. 토크나이저 pad_token 미설정으로 배치 추론 시 경고/비정상 동작
  • 채팅 모델에서 패딩이 필요하면 tokenizer.pad_token = tokenizer.eos_token 같은 설정을 고려합니다.

핵심 해법 2: FlashAttention2로 “추론을 빠르게”

FlashAttention2는 어텐션을 더 효율적으로 계산해 속도와 메모리를 동시에 개선합니다.

Transformers에서는 모델에 따라 attn_implementation 옵션으로 켤 수 있습니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "meta-llama/Llama-2-7b-chat-hf"  # 예시

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    quantization_config=bnb_config,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)

prompt = "요약: Transformers에서 로컬 LLM 최적화 방법을 알려줘."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=False,
    )

print(tokenizer.decode(out[0], skip_special_tokens=True))

FlashAttention2 적용 시 체크 포인트

  • flash-attn 설치가 필요할 수 있습니다(환경에 따라 소스 빌드가 발생).
  • GPU 아키텍처/드라이버/torch CUDA 버전의 궁합이 중요합니다.
  • 켰는데도 속도 차이가 없다면 아래를 의심하세요.
    • 실제로는 CPU 오프로딩이 많다
    • 배치/시퀀스가 너무 짧아 이득이 작다
    • torch.compile/커널 설정이 꼬였다

로딩 속도 최적화: “다운로드”와 “샤드 로딩” 줄이기

1) Hugging Face 캐시를 빠른 디스크로

예를 들어 NVMe 경로를 캐시로 지정합니다.

export HF_HOME=/mnt/nvme/hf
export TRANSFORMERS_CACHE=/mnt/nvme/hf/transformers

2) 사전 다운로드로 첫 로딩 시간 제거

서버 부팅 시 미리 받아두면, 실제 서비스 프로세스에서는 로딩만 수행합니다.

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="meta-llama/Llama-2-7b-chat-hf",
    local_dir="/mnt/nvme/models/llama2-7b",
    local_dir_use_symlinks=False,
)

그리고 로딩은 로컬 경로로:

from transformers import AutoModelForCausalLM, AutoTokenizer

path = "/mnt/nvme/models/llama2-7b"

tokenizer = AutoTokenizer.from_pretrained(path)
model = AutoModelForCausalLM.from_pretrained(path, device_map="auto")

OOM 방지 실전 팁: 생성 파라미터부터 줄여라

4bit+FlashAttention2를 켰는데도 OOM이 난다면, 대부분 KV 캐시/컨텍스트 문제입니다.

1) max_new_tokens 상한을 걸기

무제한 생성은 KV 캐시를 끝없이 키웁니다.

out = model.generate(
    **inputs,
    max_new_tokens=256,
    do_sample=True,
    temperature=0.7,
)

2) 프롬프트를 “짧게” 만드는 전처리

  • 대화 히스토리를 무한정 붙이지 말고 요약/슬라이딩 윈도우를 적용
  • RAG를 한다면 상위 문서 수를 제한

3) 배치 크기와 동시성 제한

동시 요청을 한 프로세스에서 배치로 묶으면 처리량은 좋아질 수 있지만, VRAM은 즉시 증가합니다.

디버깅: 지금 VRAM을 누가 먹는지 빠르게 확인

1) PyTorch 메모리 요약

import torch

print(torch.cuda.get_device_name(0))
print(torch.cuda.memory_summary(device=0, abbreviated=True))

2) 로딩 직후와 생성 직후를 비교

  • 로딩 직후 VRAM이 높으면 모델/오프로딩/ dtype 문제
  • 생성 후 VRAM이 급증하면 KV 캐시/컨텍스트/배치 문제

추천 조합(현실적인 기본값)

로컬 단일 GPU 기준으로 “일단 성공하고, 빠르게” 가려면 아래가 출발점으로 좋습니다.

  • bitsandbytes 4bit: nf4 + double quant + compute bfloat16
  • attn_implementation: flash_attention_2
  • device_map: auto로 시작, 목표는 GPU 상주 비율을 높이기
  • 생성 제한: max_new_tokens 하드 리밋

그리고 OOM이 계속되면, 모델 크기를 낮추는 것보다 먼저 컨텍스트/배치/동시성을 의심하는 게 보통 더 빠르게 해결됩니다.

KV 캐시까지 포함해 더 공격적으로 메모리를 줄이는 방법(예: KV 캐시 최적화, 설정 조합별 트레이드오프)은 아래 글에서 이어서 정리해두었습니다.

마무리

transformers 로컬 LLM의 “느림·OOM”은 대체로

  • 로딩 단계에서는 캐시/디스크 I/O/오프로딩
  • 추론 단계에서는 KV 캐시(컨텍스트·배치·동시성)

이 두 축에서 발생합니다.

4bit 양자화로 VRAM을 먼저 확보해 오프로딩을 줄이고, FlashAttention2로 어텐션 병목을 줄이면 체감 성능이 크게 개선됩니다. 이후에는 생성 파라미터와 운영 설정(동시성, 배치)을 조여서 안정적으로 굴리는 쪽으로 튜닝하는 것이 가장 효율적입니다.

Discuss on Twitter • View on GitHub