모델 로드

Transformers는 한 줄의 코드로 바로 사용할 수 있는 다양한 사전 학습 모델을 제공합니다. 모델 클래스와 from_pretrained() 메서드가 필요합니다.

Hugging Face Hub에 저장된 모델의 가중치와 설정을 다운로드하고 로드하려면 from_pretrained()를 호출합니다.

The from_pretrained() method loads weights stored in the safetensors file format if they’re available. Traditionally, PyTorch model weights are serialized with the pickle utility which is known to be insecure. Safetensor files are more secure and faster to load.

from transformers import AutoModelForCausalLM

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

이 가이드에서는 모델이 로드되는 방식, 모델을 로드할 수 있는 다양한 방법, 매우 큰 모델의 메모리 문제를 해결하는 방법, 그리고 커스텀 모델을 로드하는 방법을 설명합니다.

모델 및 구성

모든 모델에는 숨겨진 레이어 수, 어휘 크기, 활성화 함수 등 특정 속성이 포함된 configuration.py 파일이 있습니다. 각 레이어에서 수행되는 레이어 구조와 수학적 연산을 정의하는 modeling.py 파일도 찾을 수 있습니다. modeling.py 파일은 configuration.py의 모델 속성을 받아 해당 모델에 맞게 모델을 구성합니다. 이 시점에서 모델은 무작위 가중치를 가지며, 의미 있는 결과를 출력하도록 학습이 필요합니다.

architecture는 모델의 구조를, checkpoint는 해당 구조에 대한 모델의 가중치를 의미합니다. 예를 들어, BERT는 architecture이고 google-bert/bert-base-uncased는 checkpoint입니다. model이라는 용어는 architecture와 checkpoint를 가리키기 위해 혼용되어 사용됩니다.

로드할 수 있는 모델은 크게 두 가지 유형이 있습니다:

AutoModel 또는 LlamaModel과 같이 hidden states를 출력하는 최소한의 모델입니다.
head가 연결된 모델로, AutoModelForCausalLM 또는 LlamaForCausalLM와 같이 특정 작업을 수행하는 데 사용됩니다.

모델 클래스

사전 학습된 모델을 사용하려면 가중치를 모델에 로드해야 합니다. 이는 Hugging Face Hub 또는 로컬 디렉터리에서 가중치를 받아들이는 from_pretrained()를 호출하여 수행합니다.

모델 클래스는 두 가지가 있습니다. AutoModel 클래스와 모델별 클래스입니다.

AutoModel

모델별 클래스

대형 모델

대형 사전 학습 모델은 로딩에 많은 메모리가 필요합니다. 로딩 과정에는 다음이 포함됩니다:

무작위 가중치로 모델 생성
pretrained weights를 로드하는 것
모델에 사전 학습된 가중치를 배치하는 것

모델 가중치(랜덤 및 사전 학습)의 두 사본을 저장할 수 있는 충분한 메모리가 필요합니다. 하드웨어에 따라 이것이 불가능할 수 있습니다. 분산 학습 환경에서는 각 프로세스가 사전 학습된 모델을 로드하기 때문에 이 문제가 더욱 심각해집니다.

Transformers는 빠른 초기화, 샤드화된 체크포인트, Accelerate의 Big Model Inference 기능, 그리고 낮은 비트 데이터 타입 지원 등을 통해 이러한 메모리 관련 문제의 일부를 완화합니다.

분할된 체크포인트

save_pretrained()는 50GB를 초과하는 체크포인트를 자동으로 샤딩합니다. 이를 통해 대형 모델의 샤드 수를 적게 유지하고 파일 관리를 단순화합니다.

파라미터는 병렬로 로드되며, 피크 메모리는 모델 크기에만 의존합니다. 샤딩 전 최대 체크포인트 크기를 설정하려면 save_pretrained()에서 max_shard_size를 사용하세요.

동적 가중치 변환이 필요한 모델의 메모리 사용량은 모델 크기와 단일 변환에서 가장 큰 파라미터의 크기에 따라 달라집니다. 이는 일반적으로 mixture-of-experts(MoE) 모델에 해당하며, 메모리 사용량은 모델 크기에 한 레이어의 전문가(expert) 수를 더한 값입니다. 모델이 로드되는 방식에 대한 자세한 내용은 dynamic weight loader 가이드를 참고하세요.

save_pretrained()은 매개변수 이름과 해당 샤드 파일을 매핑하는 인덱스 파일도 생성합니다. 인덱스에는 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은 각 파라미터를 해당 샤드 파일에 매핑합니다.

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',
 ...
}

대형 모델 추론

이 기능을 사용하려면 Accelerate v0.9.0과 PyTorch v1.9.0 이상을 설치했는지 확인하세요!

from_pretrained() is supercharged with Accelerate’s Big Model Inference feature.

Big Model Inference는 PyTorch의 meta 디바이스에 모델 스켈레톤을 생성합니다. meta 디바이스는 실제 데이터를 저장하지 않고 메타데이터만 저장합니다.

pretrained 가중치가 로드될 때만 무작위로 초기화된 가중치가 생성됩니다. 이를 통해 메모리에 모델의 두 사본을 동시에 유지하는 것을 방지합니다. 최대 메모리 사용량은 모델의 크기만입니다.

디바이스 맵 설계에서 디바이스 배치에 대해 자세히 알아보세요.

Big Model Inference의 두 번째 기능은 모델 스켈레톤에서 가중치가 로드되고 배포되는 방식과 관련이 있습니다. 모델 가중치는 가장 빠른 장치(보통 GPU)부터 시작하여 모든 사용 가능한 장치에 걸쳐 배포되며, 나머지 가중치는 더 느린 장치(CPU 및 하드 드라이브)로 오프로딩됩니다.

두 기능을 함께 사용하면 대규모 사전 학습 모델의 메모리 사용량과 로딩 시간을 줄일 수 있습니다.

Big Model Inference를 활성화하려면 device_map을 "auto"(으)로 설정하세요.

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

디스크 오프로딩

모델이 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",
)

GPU 또는 CPU 메모리가 사용 불가능함을 지정하고 모델을 전체적으로 디스크로 오프로딩하도록 강제하려면 max_memory={}를 설정하세요.

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

Disk offloading은 속도를 메모리로 교환합니다. 각 forward pass마다 디스크에서 가중치를 읽는 것은 CPU 또는 GPU 메모리에서 읽는 것보다 느립니다. 완전한 disk offloading에는 accelerate>1.14.0이 필요합니다.

모델 데이터 유형

dtype 인수는 모델 가중치를 초기화할 때 사용할 PyTorch dtype을 제어합니다. 기본적으로 Transformers는 config.json의 dtype 또는 레거시 torch_dtype 값으로 가중치를 로드합니다. config.json에 해당 값이 모두 없으면, Transformers는 체크포인트의 첫 번째 부동소수점 가중치의 dtype를 사용합니다.

dtype를 전달하여 기본값을 재정의합니다.

import torch
from transformers import AutoModelForCausalLM

# specific 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의 configuration 및 modeling 클래스를 기반으로 구축되며, AutoClass API를 지원하고 from_pretrained()로 로드됩니다. 차이점은 modeling 코드가 Transformers에서 온 것이 아니라는 것입니다.

커스텀 모델을 로드할 때는 추가적인 주의를 기울이세요. Hub에는 모든 저장소에 대해 악성 코드 스캐닝이 포함되어 있지만, 악성 코드가 의도치 않게 실행되는 것을 방지하기 위해 여전히 신중해야 합니다.

커스텀 모델을 로드하려면 from_pretrained()에서 trust_remote_code=True를 설정하세요.

from transformers import AutoModelForImageClassification

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

추가 보안 조치로, 변경되었을 수 있는 모델 코드가 로드되는 것을 방지하기 위해 특정 revision에서 커스텀 모델을 로드하세요. commit hash는 모델의 commit history에서 복사할 수 있습니다.

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

자세한 내용은 모델 사용자 정의 가이드를 참조하세요.

Update on GitHub