모델 불러오기

모델 불러오기 (Loading models)

Transformers는 한 줄의 코드로 바로 쓸 수 있는 많은 사전 학습 모델을 제공해요. 모델 클래스와 [~PreTrainedModel.from_pretrained] 메서드만 있으면 돼요.

[~PreTrainedModel.from_pretrained]을 호출하면 Hugging Face Hub에 저장된 모델 가중치와 설정을 다운로드해 모델로 로드해요.

[!TIP] [~PreTrainedModel.from_pretrained]은 safetensors 파일 형식으로 저장된 가중치를 사용할 수 있다면 그걸 로드해요. 전통적으로 PyTorch 모델 가중치는 pickle 유틸리티로 직렬화되는데, 이는 안전하지 않은 것으로 알려져 있어요. safetensor 파일은 더 안전하고 로드도 더 빨라요.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", device_map="auto")

이 가이드는 모델이 어떻게 로드되는지, 모델을 로드하는 여러 방법, 정말 큰 모델의 메모리 문제를 극복하는 법, 커스텀 모델을 로드하는 법을 설명해요.

모델과 설정 (Models and configurations)

모든 모델에는 숨은 레이어 수, 어휘 크기, 활성화 함수 같은 특정 속성을 담은 configuration.py 파일이 있어요. 또 각 레이어 내부에서 일어나는 레이어와 수학 연산을 정의하는 modeling.py 파일도 찾을 수 있어요. modeling.py 파일은 configuration.py의 모델 속성을 받아 그에 맞게 모델을 만드는데, 이 시점의 모델은 의미 있는 결과를 내기 위해 훈련이 필요한 랜덤 가중치를 가진 상태예요.

[!TIP] 아키텍처는 모델의 뼈대를, 체크포인트는 주어진 아키텍처에 대한 모델 가중치를 뜻해요. 예를 들어 BERT는 아키텍처이고, google-bert/bert-base-uncased는 체크포인트예요. 모델이라는 용어는 아키텍처와 체크포인트를 함께 가리킬 때 쓰는 걸 볼 수 있을 거예요.

로드할 수 있는 모델은 크게 두 종류예요.

  1. [AutoModel]이나 [LlamaModel]처럼 hidden states를 출력하는 기본형 모델.
  2. 특정 작업을 수행하기 위해 특정 헤드가 붙은 모델. 예: [AutoModelForCausalLM], [LlamaForCausalLM].

모델 클래스 (Model classes)

사전 학습 모델을 얻으려면 모델에 가중치를 로드해야 해요. 이는 [~PreTrainedModel.from_pretrained]을 호출해서 하는데, Hugging Face Hub나 로컬 디렉토리의 가중치를 받아들여요.

AutoModel 클래스와 모델별 클래스, 두 종류의 모델 클래스가 있어요.

AutoModel: AutoModel 클래스는 아키텍처를 로드할 때 정확한 모델 클래스 이름을 몰라도 되는 편리한 방법이에요. 사용 가능한 모델이 많기 때문이죠. 설정 파일을 보고 올바른 모델 클래스를 자동으로 선택해요. 쓰려는 작업과 체크포인트만 알면 돼요.

아키텍처가 주어진 작업을 지원하기만 하면 모델이나 작업 사이를 쉽게 전환할 수 있어요.

예를 들어 같은 모델을 서로 다른 작업에 쓸 수 있어요.

from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoModelForQuestionAnswering

# 같은 API로 3가지 다른 작업에 사용
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
model = AutoModelForSequenceClassification.from_pretrained("meta-llama/Llama-2-7b-hf")
model = AutoModelForQuestionAnswering.from_pretrained("meta-llama/Llama-2-7b-hf")

다른 경우로, 한 작업에 여러 모델을 빠르게 시도해 보고 싶을 수도 있어요.

from transformers import AutoModelForCausalLM

# 같은 API로 3가지 다른 모델 로드
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("google/gemma-7b")

모델별 클래스: AutoModel 클래스는 모델별 클래스 위에 구축돼요. 특정 작업을 지원하는 모든 모델 클래스는 각자의 AutoModelFor 작업 클래스에 매핑돼요.

이미 쓰려는 모델 클래스를 알고 있다면 그 모델별 클래스를 직접 쓸 수 있어요.

from transformers import LlamaModel, LlamaForCausalLM

model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

대형 모델 (Large models)

대형 사전 학습 모델은 로드하는 데 많은 메모리를 필요로 해요. 로드 과정은 다음과 같아요.

  1. 랜덤 가중치로 모델 만들기
  2. 사전 학습 가중치 로드
  3. 사전 학습 가중치를 모델에 배치

모델 가중치 두 벌(랜덤 + 사전 학습)을 담을 메모리가 필요한데, 하드웨어에 따라 불가능할 수도 있어요. 분산 훈련 환경에서는 각 프로세스가 사전 학습 모델을 로드하므로 더 어려워요.

Transformers는 빠른 초기화, 분할 체크포인트(sharded checkpoints), Accelerate의 Big Model Inference 기능, 더 낮은 비트 데이터 타입 지원으로 이런 메모리 관련 문제 중 일부를 줄여요.

분할 체크포인트 (Sharded checkpoints)

[~PreTrainedModel.save_pretrained]은 50GB보다 큰 체크포인트를 자동으로 분할해요. 이렇게 하면 대형 모델의 shard 수를 낮게 유지하고 파일 관리를 단순화해요.

파라미터는 병렬로 로드되며 최고 메모리 사용량은 모델 크기에만 의존해요. [~PreTrainedModel.save_pretrained]의 max_shard_size를 사용해 분할 전 최대 체크포인트 크기를 설정해요.

[!NOTE] 동적 가중치 변환이 필요한 모델의 메모리 사용량은 모델 크기와 단일 변환에서 가장 큰 파라미터 크기에 따라 달라져요. 이는 일반적으로 mixture-of-experts(MoE) 모델에 적용되는데, 메모리 사용량이 모델 크기 + 한 레이어의 전문가 수가 돼요. 모델이 어떻게 로드되는지 더 알아보려면 dynamic weight loader 가이드를 참조해요.

[~PreTrainedModel.save_pretrained]은 파라미터 이름을 shard 파일에 매핑하는 인덱스 파일도 만들어요. 인덱스에는 metadata와 weight_map 두 키가 있어요.

import json
import os
import tempfile

with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir, max_shard_size="50GB")
    with open(os.path.join(tmp_dir, "model.safetensors.index.json"), "r") as f:
        index = json.load(f)

print(index.keys())

metadata는 모델 총 크기를 저장해요.

index["metadata"]
{'total_size': 28966928384}

weight_map는 각 파라미터를 자신의 shard 파일에 매핑해요.

index["weight_map"]
{'lm_head.weight': 'model-00006-of-00006.safetensors',
 'model.embed_tokens.weight': 'model-00001-of-00006.safetensors',
 'model.layers.0.input_layernorm.weight': 'model-00001-of-00006.safetensors',
 'model.layers.0.mlp.down_proj.weight': 'model-00001-of-00006.safetensors',
 ...
}

Big Model Inference

[!TIP] 이 기능을 쓰려면 Accelerate v0.9.0과 PyTorch v1.9.0 이상을 설치해야 해요!

[~PreTrainedModel.from_pretrained]은 Accelerate의 Big Model Inference 기능으로 강화됐어요.

Big Model Inference는 PyTorch meta 장치에 *모델 뼈대(skeleton)*를 만들어요. meta 장치는 실제 데이터를 저장하지 않고 메타데이터만 저장해요.

랜덤 초기화 가중치는 사전 학습 가중치가 로드될 때만 만들어져서, 메모리에 모델 두 벌을 동시에 유지하지 않아요. 최고 메모리 사용량은 모델 크기만큼이에요.

[!TIP] 장치 배치에 대해 더 알아보려면 Designing a device map을 참조해요.

Big Model Inference의 두 번째 기능은 모델 뼈대에서 가중치를 로드·배치하는 방식과 관련돼요. 모델 가중치는 가장 빠른 장치(보통 GPU)부터 시작해 사용 가능한 모든 장치에 분배되고, 남은 가중치는 느린 장치(CPU와 하드 드라이브)로 오프로드돼요.

두 기능을 합치면 큰 사전 학습 모델의 메모리 사용량과 로드 시간이 줄어요.

device_map을 "auto"로 설정해 Big Model Inference를 활성화해요.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-7b", device_map="auto")

device_map에서 레이어를 장치에 수동으로 할당할 수도 있어요. 모든 모델 파라미터를 장치에 매핑해야 하지만, 레이어 전체가 같은 장치에 있다면 레이어의 모든 서브모듈 위치를 상세히 적을 필요는 없어요.

hf_device_map 속성으로 모델이 장치들에 어떻게 분배됐는지 확인해요.

device_map = {"model.layers.1": 0, "model.layers.14": 1, "model.layers.31": "cpu", "lm_head": "disk"}
model.hf_device_map

디스크 오프로딩 (Disk offloading)

모델이 GPU와 CPU RAM을 합친 것보다 클 때, Big Model Inference는 남은 가중치를 디스크로 오프로드해요. 어떤 가중치가 "disk"에 매핑될 때마다 offload_folder를 디렉토리로 설정해요. Transformers가 오프로드된 가중치를 거기에 저장하거든요.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    offload_folder="./offload",
)

max_memory로 각 장치가 보유할 메모리 상한을 정해요. Transformers는 속도 순서대로 장치를 채우고 들어가지 않는 것은 디스크로 보내요. 아래 예시는 GPU를 10GB, CPU를 30GB로 제한해서, 남은 가중치는 offload_folder에 보관돼요.

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    max_memory={0: "10GB", "cpu": "30GB"},
    offload_folder="./offload",
)

max_memory={}를 설정하면 gpu나 cpu 메모리가 없다고 지정해서 전체 모델을 디스크로 오프로드하도록 강제해요.

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    max_memory={},
    offload_folder="./offload",
)

[!NOTE] 디스크 오프로딩은 메모리와 속도를 맞바꾸는 거예요. 매 forward pass마다 디스크에서 가중치를 읽는 것은 CPU나 GPU 메모리에서 읽는 것보다 느려요. 완전한 디스크 오프로딩에는 accelerate>1.14.0이 필요해요.

모델 데이터 타입 (Model data type)

dtype 인자는 모델 가중치를 인스턴스화하는 데 사용하는 PyTorch dtype을 제어해요. 기본적으로 Transformers는 config.json의 dtype 또는 레거시 torch_dtype 값으로 가중치를 로드해요. config.json이 둘 다 포함하지 않으면 Transformers는 체크포인트에서 첫 번째 부동소수점 가중치의 dtype을 사용해요.

기본값을 재정의하려면 특정 dtype을 넘겨요.

import torch
from transformers import AutoModelForCausalLM

# 특정 dtype
model = AutoModelForCausalLM.from_pretrained("google/gemma-3-1b-it", dtype=torch.float16)

[AutoConfig]는 처음부터 인스턴스화하는 모델에 대해서도 dtype을 받아들여요.

import torch
from transformers import AutoConfig, AutoModel

my_config = AutoConfig.from_pretrained("google/gemma-2b", dtype=torch.float16)
model = AutoModel.from_config(my_config)

커스텀 모델 (Custom models)

커스텀 모델은 Transformers의 설정·모델링 클래스 위에 구축되고, AutoClass API를 지원하며, [~PreTrainedModel.from_pretrained]으로 로드돼요. 차이점은 모델링 코드가 Transformers 것이 아니라는 점이에요.

커스텀 모델을 로드할 때는 각별히 주의해요. Hub가 모든 저장소에 멀웨어 스캔을 하기는 하지만, 실수로 악성 코드를 실행하지 않도록 여전히 주의해야 해요.

커스텀 모델을 로드하려면 [~PreTrainedModel.from_pretrained]에서 trust_remote_code=True를 설정해요.

from transformers import AutoModelForImageClassification

model = AutoModelForImageClassification.from_pretrained("sgugger/custom-resnet50d", trust_remote_code=True)

추가 보안 계층으로, 변경될 수 있는 모델 코드를 로드하지 않도록 특정 리비전에서 커스텀 모델을 로드해요. 커밋 해시는 모델의 커밋 히스토리에서 복사할 수 있어요.

commit_hash = "ed94a7c6247d8aedce4647f00f20de6875b5b292"
model = AutoModelForImageClassification.from_pretrained(
    "sgugger/custom-resnet50d", trust_remote_code=True, revision=commit_hash
)

더 자세한 내용은 Customize models 가이드를 참조해요.

출처: Hugging Face Transformers — Loading models

더 알아보기 (Learn more)