Tải mô hình

Transformers cung cấp nhiều mô hình được huấn luyện trước, sẵn sàng sử dụng chỉ với một dòng mã. Bạn cần một lớp mô hình và phương thức from_pretrained().

Gọi from_pretrained() để tải xuống và nạp trọng số cùng cấu hình của một mô hình được lưu trữ trên Hub của Hugging Face.

Phương thức from_pretrained() sẽ nạp trọng số được lưu trữ ở định dạng tệp safetensors nếu có sẵn. Truyền thống, trọng số mô hình PyTorch được chuỗi hóa bằng tiện ích pickle, vốn được biết đến là không an toàn. Tệp Safetensor an toàn hơn và nạp nhanh hơn.

from transformers import AutoModelForCausalLM

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

Hướng dẫn này giải thích cách các mô hình được nạp, các phương pháp khác nhau để nạp một mô hình, cách khắc phục các vấn đề về bộ nhớ đối với những mô hình rất lớn, và cách nạp các mô hình tùy chỉnh.

Mô hình và cấu hình

Tất cả các model đều có một tệp configuration.py chứa các thuộc tính cụ thể như số lớp ẩn, kích thước từ vựng, hàm kích hoạt, và nhiều hơn nữa. Bạn cũng sẽ tìm thấy tệp modeling.py định nghĩa các lớp và các phép toán học diễn ra bên trong từng lớp. Tệp modeling.py nhận các thuộc tính model trong configuration.py và xây dựng model tương ứng. Ở bước này, bạn đã có một model với các trọng số ngẫu nhiên, cần được huấn luyện để đưa ra các kết quả có ý nghĩa.

Kiến trúc (architecture) chỉ bộ khung của mô hình, còn checkpoint chỉ trọng số của mô hình cho một kiến trúc nhất định. Ví dụ, BERT là một kiến trúc trong khi google-bert/bert-base-uncased là một checkpoint. Bạn sẽ thấy thuật ngữ mô hình được dùng thay thế cho cả kiến trúc và checkpoint.

Có hai loại mô hình chung mà bạn có thể tải:

Một model tối giản, chẳng hạn như AutoModel hoặc LlamaModel, có chức năng xuất ra hidden states.
Một model có gắn một head cụ thể, chẳng hạn như AutoModelForCausalLM hoặc LlamaForCausalLM, để thực hiện các tác vụ cụ thể.

Các lớp mô hình

Để có một mô hình được huấn luyện trước, bạn cần nạp trọng số vào mô hình. Việc này được thực hiện bằng cách gọi from_pretrained(), hàm này chấp nhận trọng số từ Hugging Face Hub hoặc một thư mục cục bộ.

Có hai lớp mô hình, đó là lớp AutoModel và một lớp dành riêng cho mô hình.

AutoModel

lớp cụ thể cho mô hình

Các mô hình lớn

Các mô hình được huấn luyện trước có kích thước lớn đòi hỏi nhiều bộ nhớ để tải. Quá trình tải bao gồm:

tạo một mô hình với trọng số ngẫu nhiên
tải trọng số được huấn luyện trước
đặt các trọng số được huấn luyện trước lên mô hình

Bạn cần đủ bộ nhớ để chứa hai bản sao của trọng số mô hình (ngẫu nhiên và được huấn luyện trước), điều này có thể không khả thi tùy thuộc vào phần cứng của bạn. Trong môi trường huấn luyện phân tán, điều này còn khó khăn hơn vì mỗi tiến trình đều tải một mô hình được huấn luyện trước.

Transformers giảm bớt một số thách thức liên quan đến bộ nhớ này bằng cách khởi tạo nhanh, checkpoint phân mảnh, tính năng Big Model Inference của Accelerate, và hỗ trợ các kiểu dữ liệu bit thấp hơn.

Điểm kiểm tra phân mảnh

save_pretrained() tự động chia nhỏ các checkpoint có kích thước lớn hơn 50GB. Điều này giúp giảm số lượng shard cho các mô hình lớn và đơn giản hóa việc quản lý tệp.

Các tham số được nạp song song và bộ nhớ đỉnh chỉ phụ thuộc vào kích thước mô hình. Sử dụng max_shard_size trong save_pretrained() để đặt kích thước tối đa của checkpoint trước khi chia nhỏ.

Lượng bộ nhớ sử dụng cho các mô hình yêu cầu chuyển đổi trọng số động phụ thuộc vào kích thước mô hình và kích thước của các tham số lớn nhất trong một lần chuyển đổi. Điều này thường áp dụng cho các mô hình mixture-of-experts (MoE), trong đó lượng bộ nhớ sử dụng bằng kích thước mô hình cộng với số lượng chuyên gia trên một lớp. Tham khảo hướng dẫn dynamic weight loader để tìm hiểu thêm về cách các mô hình được nạp.

save_pretrained() cũng tạo ra một tệp chỉ mục ánh xạ tên tham số đến tệp phân mảnh tương ứng. Tệp chỉ mục chứa hai khóa, metadata và 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 lưu trữ kích thước tổng của mô hình.

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

weight_map ánh xạ mỗi tham số đến tệp shard tương ứng của nó.

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

Suy luận Mô hình Lớn

Hãy đảm bảo bạn đã cài đặt Accelerate v0.9.0 và PyTorch v1.9.0 hoặc các phiên bản mới hơn để sử dụng tính năng này!

from_pretrained() được tăng cường mạnh mẽ nhờ tính năng Big Model Inference của Accelerate.

Big Model Inference tạo khung mô hình trên thiết bị meta của PyTorch. Thiết bị meta không lưu trữ bất kỳ dữ liệu thực nào, chỉ lưu trữ metadata.

Trọng số được khởi tạo ngẫu nhiên chỉ được tạo ra khi tải trọng số được huấn luyện trước, nhằm tránh việc duy trì hai bản sao của mô hình trong bộ nhớ cùng một lúc. Mức sử dụng bộ nhớ tối đa chỉ bằng kích thước của mô hình.

Tìm hiểu thêm về việc đặt thiết bị trong Thiết kế bản đồ thiết bị.

Tính năng thứ hai của Big Model Inference liên quan đến cách tải và phân phối trọng số trong khung mô hình. Trọng số mô hình được phân phối trên tất cả các thiết bị khả dụng, bắt đầu từ thiết bị nhanh nhất (thường là GPU) và sau đó chuyển các trọng số còn lại sang các thiết bị chậm hơn (CPU và ổ cứng).

Việc kết hợp cả hai tính năng giúp giảm mức sử dụng bộ nhớ và thời gian tải cho các mô hình được huấn luyện trước có kích thước lớn.

Đặt device_map thành "auto" để bật Big Model Inference.

from transformers import AutoModelForCausalLM

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

Bạn cũng có thể gán thủ công các lớp vào một thiết bị trong device_map. Nó phải ánh xạ tất cả các tham số của mô hình đến một thiết bị, nhưng bạn không cần chỉ rõ vị trí của tất cả các mô-đun con của một lớp nếu toàn bộ lớp đó nằm trên cùng một thiết bị.

Truy cập thuộc tính hf_device_map để xem mô hình được phân phối trên các thiết bị như thế nào.

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

Tải xuống ổ đĩa

Khi một model quá lớn so với tổng dung lượng của GPU và RAM CPU, Big Model Inference sẽ offload các weights còn lại ra disk. Hãy đặt offload_folder thành một thư mục bất cứ khi nào có weight nào được map tới "disk", vì đó là nơi Transformers lưu trữ các weights đã được offload.

from transformers import AutoModelForCausalLM

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

Sử dụng max_memory để giới hạn lượng bộ nhớ mà mỗi thiết bị có thể giữ. Transformers sẽ lấp đầy các thiết bị theo thứ tự tốc độ và chuyển phần không vừa sang ổ đĩa. Ví dụ bên dưới giới hạn GPU ở mức 10GB và CPU ở mức 30GB, do đó các trọng số còn lại sẽ được giữ trong offload_folder.

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

Đặt max_memory={} để chỉ định rằng không có bộ nhớ gpu hoặc cpu khả dụng và để buộc việc offload toàn bộ model ra disk.

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

Disk offloading đánh đổi bộ nhớ lấy tốc độ. Việc đọc trọng số từ đĩa trong mỗi bước truyền thẳng chậm hơn so với việc đọc từ bộ nhớ CPU hoặc GPU. Disk offloading hoàn toàn yêu cầu accelerate>1.14.0

Loại dữ liệu của mô hình

Tham số dtype kiểm soát dtype của PyTorch được sử dụng để khởi tạo trọng số mô hình. Mặc định, Transformers nạp trọng số với giá trị dtype hoặc giá trị cũ torch_dtype từ config.json. Nếu config.json không chứa giá trị nào trong hai giá trị trên, Transformers sẽ sử dụng dtype của trọng số số thực đầu tiên trong checkpoint.

Ghi đè mặc định bằng cách truyền một dtype cụ thể.

import torch
from transformers import AutoModelForCausalLM

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

AutoConfig cũng chấp nhận dtype cho các mô hình được khởi tạo từ đầu.

import torch
from transformers import AutoConfig, AutoModel

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

Mô hình tùy chỉnh

Mô hình tùy chỉnh được xây dựng dựa trên các lớp cấu hình và mô hình hóa của Transformers, hỗ trợ API AutoClass và được nạp bằng from_pretrained(). Điểm khác biệt là mã mô hình hóa không đến từ Transformers.

Hãy đặc biệt cẩn thận khi tải một model tùy chỉnh. Mặc dù Hub bao gồm quét phần mềm độc hại cho mọi repository, bạn vẫn nên thận trọng để tránh vô tình thực thi mã độc.

Đặt trust_remote_code=True trong from_pretrained() để tải một model tùy chỉnh.

from transformers import AutoModelForImageClassification

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

Để tăng thêm một lớp bảo mật, hãy tải một model tùy chỉnh từ một revision cụ thể để tránh tải mã model có thể đã bị thay đổi. Hash của commit có thể được sao chép từ lịch sử commit của các model.

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

Tham khảo hướng dẫn Tùy chỉnh mô hình để biết thêm thông tin.

Update on GitHub