Hugging Face's logo

Accelerate documentation

Nạp các mô hình lớn vào bộ nhớ

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.15.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Nạp các mô hình lớn vào bộ nhớ

Khi tải một mô hình được huấn luyện trước trong PyTorch, quy trình thông thường trông như sau:

import torch

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

Nói một cách đơn giản, các bước đó là:

  1. Tạo mô hình với các trọng số được khởi tạo ngẫu nhiên
  2. Tải trọng số mô hình (thường là một từ điển được gọi là state dict) từ ổ đĩa
  3. Tải các trọng số đó vào bên trong mô hình

Mặc dù cách này hoạt động rất tốt với các model có kích thước thông thường, quy trình này có một số hạn chế rõ ràng khi chúng ta làm việc với một model khổng lồ: ở bước 1, chúng ta nạp một phiên bản đầy đủ của model vào RAM và dành một khoảng thời gian để khởi tạo ngẫu nhiên các trọng số (những trọng số này sẽ bị loại bỏ ở bước 3). Ở bước 2, chúng ta nạp một phiên bản đầy đủ khác của model vào RAM, với các trọng số đã được huấn luyện trước. Nếu bạn đang nạp một model có 6 tỷ tham số, điều này có nghĩa là bạn sẽ cần 24GB RAM cho mỗi bản sao của model, tức là 48GB tổng cộng (một nửa trong số đó để nạp model ở FP16).

API này khá mới và vẫn đang ở giai đoạn thử nghiệm. Mặc dù chúng tôi nỗ lực cung cấp một API ổn định, nhưng một số phần nhỏ của API công khai có thể sẽ thay đổi trong tương lai.

Quy trình hoạt động: Tổng quan nhanh

Cách quy trình hoạt động: Làm việc với mã

Tạo một mô hình rỗng

Công cụ đầu tiên mà Accelerate giới thiệu để hỗ trợ các mô hình lớn là trình quản lý ngữ cảnh init_empty_weights(), giúp bạn khởi tạo một mô hình mà không sử dụng bất kỳ bộ nhớ RAM nào, nhờ đó bước 1 có thể được thực hiện trên các mô hình có kích thước bất kỳ. Dưới đây là cách nó hoạt động:

from accelerate import init_empty_weights

with init_empty_weights():
    my_model = ModelClass(...)

Ví dụ:

with init_empty_weights():
    model = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])

khởi tạo một mô hình trống với hơn 100 tỷ tham số một chút. Đằng sau, điều này dựa vào thiết bị meta được giới thiệu trong PyTorch 1.9. Trong quá trình khởi tạo dưới trình quản lý ngữ cảnh, mỗi khi một tham số được tạo ra, nó sẽ ngay lập tức được chuyển sang thiết bị đó.

Bạn không thể di chuyển một model được khởi tạo theo cách này từ CPU hoặc một thiết bị khác một cách trực tiếp, vì nó không chứa bất kỳ dữ liệu nào. Ngoài ra, rất có thể một bước forward với model rỗng đó sẽ thất bại, vì không phải tất cả các phép toán đều được hỗ trợ trên meta device.

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

Có thể mô hình của bạn quá lớn đến mức ngay cả một bản sao duy nhất cũng không thể chứa vừa trong RAM. Điều đó không có nghĩa là nó không thể được nạp: nếu bạn có một hoặc nhiều GPU, đây là bộ nhớ bổ sung có sẵn để lưu trữ mô hình của bạn. Trong trường hợp này, tốt hơn nếu checkpoint của bạn được chia thành nhiều tệp nhỏ hơn mà chúng tôi gọi là các shard của checkpoint.

Accelerate sẽ xử lý các checkpoint dạng sharded miễn là bạn tuân thủ định dạng sau: checkpoint của bạn phải nằm trong một thư mục, với một số tệp chứa các state dict một phần, và phải có một index ở định dạng JSON chứa một từ điển ánh xạ tên tham số đến tệp chứa trọng số của chúng. Bạn có thể dễ dàng chia nhỏ (shard) mô hình của mình bằng save_model(). Ví dụ, chúng ta có thể có một thư mục chứa:

first_state_dict.bin
index.json
second_state_dict.bin

với index.json là tệp sau:

{
  "linear1.weight": "first_state_dict.bin",
  "linear1.bias": "first_state_dict.bin",
  "linear2.weight": "second_state_dict.bin",
  "linear2.bias": "second_state_dict.bin"
}

và first_state_dict.bin chứa trọng số cho "linear1.weight" và "linear1.bias", second_state_dict.bin chứa các trọng số cho "linear2.weight" và "linear2.bias"

Tải trọng số

Công cụ thứ hai mà Accelerate giới thiệu là hàm load_checkpoint_and_dispatch(), cho phép bạn nạp checkpoint vào bên trong model trống của mình. Hàm này hỗ trợ cả checkpoint đầy đủ (một tệp duy nhất chứa toàn bộ state dict) lẫn checkpoint dạng sharded. Nó cũng sẽ tự động phân phối các trọng số đó trên các thiết bị bạn có sẵn (GPU, RAM CPU), do đó nếu bạn đang nạp một checkpoint dạng sharded, mức sử dụng RAM tối đa sẽ bằng kích thước của shard lớn nhất.

Nếu bạn muốn sử dụng suy luận mô hình lớn với các mô hình Transformers, hãy xem tài liệu này.

Dưới đây là cách chúng ta có thể sử dụng điều này để tải mô hình GPT2-1.5B.

Hãy tải xuống phiên bản chia mảnh của mô hình này.

pip install huggingface_hub
from huggingface_hub import snapshot_download
checkpoint = "marcsun13/gpt2-xl-linear-sharded"
weights_location = snapshot_download(repo_id=checkpoint)

Để khởi tạo mô hình, chúng ta sẽ sử dụng thư viện minGPT.

git clone https://github.com/karpathy/minGPT.git
pip install minGPT/
from accelerate import init_empty_weights
from mingpt.model import GPT

model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024

with init_empty_weights():
    model = GPT(model_config)

Sau đó, tải checkpoint vừa tải xuống bằng cách sử dụng:

from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map="auto", no_split_module_classes=['Block']
)

Bằng cách truyền device_map="auto", chúng ta cho Accelerate biết để tự động xác định vị trí đặt từng lớp của mô hình dựa trên các tài nguyên khả dụng:

  • đầu tiên, chúng ta sử dụng toàn bộ dung lượng trống có sẵn trên GPU
  • nếu chúng ta vẫn cần thêm dung lượng, chúng ta sẽ lưu trữ các trọng số còn lại trên CPU
  • nếu không đủ RAM, chúng ta sẽ lưu trữ các trọng số còn lại trên ổ cứng dưới dạng các tensor ánh xạ bộ nhớ

no_split_module_classes

Tham số này sẽ chỉ ra rằng một số module có tên "Block" không nên được chia tách giữa các thiết bị khác nhau. Bạn nên đặt ở đây tất cả các block bao gồm một kết nối residual ở dạng nào đó.

The device_map

Bạn có thể xem device_map mà Accelerate đã chọn bằng cách truy cập thuộc tính hf_device_map của mô hình của bạn:

model.hf_device_map
{'transformer.wte': 0,
 'transformer.wpe': 0,
 'transformer.drop': 0,
 'transformer.h.0': 0,
 ...
 'transformer.h.21': 0, 
 'transformer.h.22': 1, 
 'transformer.h.23': 1, 
 'transformer.h.24': 1,
 ...
 'transformer.h.47': 1, 
 'transformer.ln_f': 1, 
 'lm_head': 1}

Hoàn toàn có thể tự tạo bản đồ thiết bị (device map) cho các lớp (layers) và chỉ định thiết bị GPU cần dùng (một số), "cpu" hoặc "disk", sau đó truyền vào:

device_map = {
    "transformer.wte": "cpu",
    "transformer.wpe": 0,
    "transformer.drop": "cpu",
    "transformer.h.0": "disk"
}

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map=device_map
)

Chạy mô hình

Bây giờ, sau khi thực hiện các bước này, mô hình của chúng ta trải dài trên nhiều thiết bị, và có thể cả ổ cứng. Nhưng nó vẫn có thể được sử dụng như một mô hình PyTorch thông thường:

from mingpt.bpe import BPETokenizer
tokenizer = BPETokenizer()
inputs = tokenizer("Hello, my name is").to(0)

outputs = model.generate(x1, max_new_tokens=10, do_sample=False)[0]
tokenizer.decode(outputs.cpu().squeeze())

Phía sau, Accelerate đã thêm các hook vào mô hình, để:

  • tại mỗi lớp, các đầu vào được đặt trên thiết bị phù hợp (nên ngay cả khi mô hình của bạn được phân bổ trên nhiều GPU, nó vẫn hoạt động)
  • đối với các trọng số được chuyển xuống CPU, chúng sẽ được đưa lên GPU ngay trước bước truyền thẳng và được dọn dẹp ngay sau đó
  • đối với các trọng số được chuyển xuống ổ cứng, chúng được nạp vào RAM sau đó đưa lên GPU ngay trước bước truyền thẳng và được dọn dẹp ngay sau đó

Nhờ đó, mô hình của bạn vẫn có thể chạy suy luận ngay cả khi không vừa với bộ nhớ của một GPU hoặc RAM CPU!

Tính năng này chỉ hỗ trợ suy luận (inference) mô hình của bạn, không hỗ trợ huấn luyện. Phần lớn các phép tính được thực hiện bên trong các trình quản lý ngữ cảnh torch.no_grad() để tránh tốn bộ nhớ GPU cho các kích hoạt trung gian.

Thiết kế bản đồ thiết bị

Bạn có thể để Accelerate xử lý việc tính toán bản đồ thiết bị bằng cách đặt device_map thành một trong các tùy chọn được hỗ trợ ("auto", "balanced", "balanced_low_0", "sequential") hoặc tự tạo một bản đồ nếu bạn muốn kiểm soát nhiều hơn về vị trí của từng lớp.

Bạn có thể suy ra tất cả các kích thước của mô hình (và do đó tính toán device_map) trên một mô hình đang nằm trên meta device.

Tất cả các tùy chọn sẽ cho ra cùng một kết quả khi bạn không có đủ bộ nhớ GPU để chứa toàn bộ mô hình (tức là đặt mọi thứ có thể lên GPU, sau đó chuyển trọng số sang CPU hoặc thậm chí sang ổ đĩa nếu RAM không đủ).

Khi bạn có nhiều bộ nhớ GPU hơn kích thước mô hình, đây là sự khác biệt giữa các tùy chọn:

  • "auto" và "balanced" chia đều mô hình trên tất cả các GPU khả dụng, cho phép bạn sử dụng kích thước lô lớn hơn 1.
  • "balanced_low_0" chia đều mô hình trên tất cả các GPU ngoại trừ GPU đầu tiên, và chỉ đặt lên GPU 0 những phần không vừa trên các GPU còn lại. Tùy chọn này rất hữu ích khi bạn cần sử dụng GPU 0 để xử lý một số đầu ra, chẳng hạn như khi sử dụng hàm generate cho các mô hình Transformers
  • "sequential" sẽ nạp vừa đủ vào GPU 0, sau đó chuyển sang GPU 1 và tiếp tục như vậy (do đó sẽ không sử dụng các GPU cuối nếu không cần thiết).

Các tùy chọn "auto" và "balanced" hiện cho ra kết quả giống nhau, nhưng hành vi của "auto" có thể thay đổi trong tương lai nếu chúng tôi tìm được một chiến lược hợp lý hơn, trong khi "balanced" sẽ giữ nguyên ổn định.

Lưu ý đầu tiên là bạn có thể giới hạn bộ nhớ sử dụng trên mỗi GPU bằng cách sử dụng đối số max_memory (khả dụng trong infer_auto_device_map() và trong tất cả các hàm sử dụng nó). Khi thiết lập max_memory, bạn nên truyền một từ điển chứa các định danh GPU (ví dụ 0, 1 v.v.) và khóa "cpu" cho bộ nhớ RAM tối đa bạn muốn sử dụng cho CPU offload. Các giá trị có thể là một số nguyên (theo byte) hoặc một chuỗi đại diện cho một số kèm theo đơn vị, chẳng hạn như "10GiB" hoặc "10GB".

Dưới đây là một ví dụ trong đó chúng ta không muốn sử dụng quá 10GiB trên mỗi trong hai GPU và không quá 30GiB RAM CPU cho trọng số mô hình:

from accelerate import infer_auto_device_map

device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})

Khi một lần cấp phát bộ nhớ đầu tiên xảy ra trong PyTorch, nó sẽ nạp các kernel CUDA, chiếm khoảng 1-2GB bộ nhớ tùy thuộc vào GPU. Do đó, bạn luôn có ít bộ nhớ khả dụng hơn so với kích thước thực tế của GPU. Để xem lượng bộ nhớ thực sự được sử dụng, hãy chạy torch.ones(1).cuda() và kiểm tra mức sử dụng bộ nhớ.

Do đó, khi bạn tạo memory maps với max_memory, hãy đảm bảo điều chỉnh bộ nhớ khả dụng cho phù hợp để tránh lỗi out-of-memory.

Ngoài ra, nếu bạn thực hiện một số thao tác bổ sung với kết quả đầu ra mà không đưa chúng trở lại CPU (ví dụ bên trong phương thức generate của Transformers) và nếu bạn đã đặt các đầu vào trên GPU, GPU đó sẽ tiêu thụ nhiều bộ nhớ hơn các GPU khác (Accelerate luôn đặt kết quả đầu ra trở lại thiết bị của đầu vào). Do đó, nếu bạn muốn tối ưu hóa kích thước lô (batch size) tối đa và bạn có nhiều GPU, hãy cấp ít bộ nhớ hơn cho GPU đầu tiên. Ví dụ, với BLOOM-176B trên cấu hình 8x80 A100, bản đồ gần lý tưởng nhất là:

max_memory = {0: "30GIB", 1: "46GIB", 2: "46GIB", 3: "46GIB", 4: "46GIB", 5: "46GIB", 6: "46GIB", 7: "46GIB"}

như bạn có thể thấy, chúng tôi đã cấp cho 7 GPU còn lại nhiều bộ nhớ hơn GPU 0 khoảng 50%.

Nếu bạn chọn tự thiết kế hoàn toàn device_map, nó phải là một từ điển với các khóa là tên module của mô hình và các giá trị là định danh thiết bị hợp lệ (ví dụ một số nguyên cho các GPU) hoặc "cpu" cho CPU offload, "disk" cho disk offload. Các khóa cần bao phủ toàn bộ mô hình, sau đó bạn có thể định nghĩa bản đồ thiết bị của mình theo ý muốn: ví dụ, nếu mô hình của bạn có hai khối (gọi là block1 và block2) mỗi khối chứa ba lớp tuyến tính (gọi là linear1, linear2 và linear3), một bản đồ thiết bị hợp lệ có thể là:

device_map = {"block1": 0, "block2": 1}

một ví dụ khác hợp lệ có thể là:

device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}

Mặt khác, cái này không hợp lệ vì nó không bao phủ mọi tham số của mô hình:

device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}

Để đạt hiệu quả cao nhất, hãy đảm bảo rằng bản đồ thiết bị của bạn đặt các tham số lên các GPU theo trình tự (ví dụ: không đặt một trong các trọng số đầu tiên lên GPU 0, sau đó đặt các trọng số lên GPU 1 và đưa trọng số cuối cùng trở lại GPU 0) để tránh việc phải chuyển dữ liệu giữa các GPU quá nhiều lần.

Chỉ offload CPU

Nếu bạn muốn offload mô hình của mình lên CPU, bạn có thể sử dụng cpu_offload(). Kết quả là, tất cả các tham số của mô hình sẽ được offload và chỉ giữ lại một bản sao của state dict của mô hình. Trong quá trình forward pass, các tham số sẽ được trích xuất từ state dict đó và đặt lên thiết bị thực thi, sau đó được truyền đi khi cần thiết, rồi lại được offload.

cpu_offload(model, execution_device)

Bạn cũng có thể sử dụng cpu_offload_with_hook(). Hàm này sẽ offload mô hình lên CPU và đưa nó trở lại thiết bị thực thi khi được thực thi. Sự khác biệt so với cpu_offload() là mô hình sẽ ở lại trên thiết bị thực thi sau khi forward và chỉ được offload lại khi phương thức offload của hook được trả về được gọi. Hơn nữa, cpu_offload_with_hook() có hiệu suất cao hơn nhưng tiết kiệm bộ nhớ kém hơn. Nó hữu ích cho các pipeline chạy mô hình trong một vòng lặp:

model_1, hook_1 = cpu_offload_with_hook(model_1, execution_device)
model_2, hook_2 = cpu_offload_with_hook(model_2, execution_device, prev_module_hook=hook_1)
model_3, hook_3 = cpu_offload_with_hook(model_3, execution_device, prev_module_hook=hook_2)

hid_1 = model_1(input)
for i in range(50):
    # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
    hid_2 = model_2(hid_1)
# model2 is offloaded to the CPU just before this forward.
hid_3 = model_3(hid_3)

# For model3, you need to manually call the hook offload method.
hook_3.offload()

Chỉ offload xuống đĩa

Để thực hiện disk offload, bạn có thể sử dụng disk_offload(). Kết quả là, tất cả các tham số của mô hình sẽ được offload dưới dạng mảng memory-mapped trong một thư mục được chỉ định. Trong quá trình forward pass, các tham số sẽ được truy cập từ thư mục đó và được đưa lên thiết bị thực thi được truyền vào khi cần, sau đó lại được offload.

disk_offload(model, offload_dir, execution_device)

Hạn chế và phát triển thêm

Chúng tôi nhận thức được các hạn chế hiện tại của API:

  • infer_auto_device_map() (hoặc device_map="auto" trong load_checkpoint_and_dispatch()) cố gắng tối đa hóa bộ nhớ RAM của GPU và CPU mà nó nhìn thấy khả dụng khi bạn thực thi. Trong khi PyTorch rất giỏi trong việc quản lý bộ nhớ RAM của GPU một cách hiệu quả (và trả lại bộ nhớ khi không cần), điều này không hoàn toàn đúng với Python và bộ nhớ RAM của CPU. Do đó, bản đồ thiết bị được tính toán tự động có thể gây quá tải cho CPU. Hãy di chuyển một vài module sang thiết bị đĩa nếu bạn gặp sự cố do thiếu RAM.
  • infer_auto_device_map() (hoặc device_map="auto" trong load_checkpoint_and_dispatch()) gán thiết bị theo thứ tự (để tránh việc di chuyển dữ liệu qua lại) nên nếu lớp đầu tiên của bạn lớn hơn kích thước GPU mà bạn có, mọi thứ sẽ bị đặt lên CPU/Disk.
  • load_checkpoint_and_dispatch() và load_checkpoint_in_model() hiện tại không thực hiện bất kỳ kiểm tra nào về tính chính xác của state dict so với model của bạn (điều này sẽ được khắc phục trong một phiên bản tương lai), do đó bạn có thể gặp một số lỗi kỳ lạ nếu cố gắng tải checkpoint có các key không khớp hoặc bị thiếu.
  • Song song mô hình được sử dụng khi mô hình của bạn được chia trên nhiều GPU là đơn giản và chưa được tối ưu, có nghĩa là chỉ có một GPU hoạt động tại một thời điểm trong khi các GPU khác ở trạng thái chờ.
  • Khi trọng số được chuyển xuống CPU/ổ cứng, hiện tại chưa có cơ chế lấy trước (chúng tôi sẽ làm việc cho các phiên bản tương lai), điều này có nghĩa là trọng số chỉ được đưa lên GPU khi cần thiết chứ không phải trước đó.
  • Việc offload xuống ổ cứng có thể rất chậm nếu phần cứng bạn sử dụng không có kết nối truyền thông nhanh giữa ổ đĩa và CPU (ví dụ như NVMes).
Update on GitHub

←Cơ chế bên trong của Accelerate So sánh hiệu năng giữa các cấu hình phân tán→