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

Một trong những bước tiến lớn nhất mà Accelerate mang lại là Big Model Inference, cho phép bạn thực hiện suy luận với các mô hình không vừa hoàn toàn trên card đồ họa của bạn.

Tutorial này sẽ hướng dẫn bạn cách sử dụng Big Model Inference trong Accelerate và hệ sinh thái Hugging Face.

Accelerate

Quy trình điển hình để tải một mô hình PyTorch được minh họa bên dưới. ModelClass là một mô hình vượt quá bộ nhớ GPU của thiết bị của bạn (mps hoặc cuda hoặc xpu).

import torch

my_model = ModelClass(...)
state_dict = torch.load(checkpoint_file)
my_model.load_state_dict(state_dict)

Với Big Model Inference, bước đầu tiên là khởi tạo một khung xương trống của mô hình bằng trình quản lý ngữ cảnh init_empty_weights. Thao tác này không yêu cầu bộ nhớ vì my_model là “không có tham số”.

from accelerate import init_empty_weights
with init_empty_weights():
    my_model = ModelClass(...)

Tiếp theo, các trọng số được nạp vào mô hình để suy luận.

Phương thức load_checkpoint_and_dispatch() tải checkpoint vào mô hình trống của bạn và phân phối trọng số cho từng lớp trên tất cả các thiết bị khả dụng, bắt đầu với các thiết bị nhanh nhất (GPU, MPS, XPU, NPU, MLU, SDAA, MUSA) trước khi chuyển sang các thiết bị chậm hơn (CPU và ổ cứng).

Việc đặt device_map="auto" sẽ tự động lấp đầy toàn bộ dung lượng trống trên GPU trước, sau đó là CPU, và cuối cùng là ổ cứng (tùy chọn chậm nhất tuyệt đối) nếu vẫn còn thiếu bộ nhớ.

Tham khảo hướng dẫn Thiết kế bản đồ thiết bị để biết thêm chi tiết về cách thiết kế bản đồ thiết bị của riêng bạn.

from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=checkpoint_file, device_map="auto"
)

Nếu có một số “khối” lớp nhất định không nên tách, hãy truyền chúng vào no_split_module_classes (xem ở đây để biết thêm chi tiết).

Trọng số của một mô hình cũng có thể được chia nhỏ thành nhiều checkpoint để tiết kiệm bộ nhớ, chẳng hạn khi state_dict không vừa trong bộ nhớ (xem ở đây để biết thêm chi tiết).

Bây giờ mô hình đã được phân bổ hoàn toàn, bạn có thể thực hiện suy luận.

input = torch.randn(2,3)
device_type = next(iter(model.parameters())).device.type
input = input.to(device_type)
output = model(input)

Mỗi lần một đầu vào được truyền qua một lớp, nó được gửi từ CPU sang GPU (hoặc từ ổ đĩa sang CPU rồi sang GPU), đầu ra được tính toán, và lớp đó được gỡ khỏi GPU để di chuyển ngược xuống phía dưới. Mặc dù điều này làm tăng một chút chi phí xử lý cho suy luận, nhưng nó cho phép bạn chạy bất kỳ mô hình nào có kích thước trên hệ thống của mình, miễn là lớp lớn nhất vừa đủ trên GPU của bạn.

Có thể sử dụng nhiều GPU, hay “model parallelism”, nhưng tại mỗi thời điểm chỉ có một GPU hoạt động. Điều này buộc GPU phải chờ GPU trước đó gửi kết quả đầu ra. Bạn nên chạy script bằng Python thông thường thay vì các công cụ khác như torchrun và accelerate launch.

Xem một ví dụ đầy đủ về suy luận mô hình lớn bên dưới.

import torch
from accelerate import init_empty_weights, load_checkpoint_and_dispatch

with init_empty_weights():
    model = MyModel(...)

model = load_checkpoint_and_dispatch(
    model, checkpoint=checkpoint_file, device_map="auto"
)

input = torch.randn(2,3)
device_type = next(iter(model.parameters())).device.type
input = input.to(device_type)
output = model(input)

Hệ sinh thái Hugging Face

Các thư viện khác trong hệ sinh thái Hugging Face, chẳng hạn như Transformers hoặc Diffusers, hỗ trợ Big Model Inference trong các bộ khởi tạo from_pretrained của chúng.

Bạn chỉ cần thêm device_map="auto" vào from_pretrained để bật Big Model Inference.

Ví dụ, nạp mô hình 11 tỷ tham số Big Sciences T0pp bằng Big Model Inference.

from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto")

Sau khi tải mô hình, các bước khởi tạo rỗng và điều phối thông minh như trước đó sẽ được thực thi và mô hình đã sẵn sàng hoàn toàn để tận dụng tất cả các tài nguyên trên máy của bạn. Thông qua các bộ khởi tạo này, bạn cũng có thể tiết kiệm thêm bộ nhớ bằng cách chỉ định tham số torch_dtype để tải mô hình ở độ chính xác thấp hơn.

from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto", torch_dtype=torch.float16)

Các bước tiếp theo

Để có giải thích chi tiết hơn về Big Model Inference, hãy chắc chắn kiểm tra hướng dẫn khái niệm!