Memuat model besar ke dalam memori

Saat memuat model terlatih sebelumnya di PyTorch, alur kerja yang umum terlihat seperti ini:

import torch

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

Secara sederhana, langkah-langkah tersebut adalah:

1. Buat model dengan bobot yang diinisialisasi secara acak
2. Muat bobot model (dalam kamus yang biasanya disebut state dict) dari disk
3. Muat bobot tersebut di dalam model

Meskipun pendekatan ini bekerja sangat baik untuk model berukuran reguler, alur kerja ini memiliki beberapa keterbatasan yang jelas ketika kita berurusan dengan model yang sangat besar: pada langkah 1, kami memuat versi lengkap model ke RAM dan menghabiskan waktu untuk menginisialisasi bobot secara acak (yang akan dibuang pada langkah 3). Pada langkah 2, kami memuat versi lengkap model lainnya ke RAM, dengan bobot terlatih sebelumnya. Jika Anda memuat model dengan 6 miliar parameter, ini berarti Anda akan membutuhkan 24GB RAM untuk setiap salinan model, sehingga total 48GB (setengahnya untuk memuat model dalam FP16).

API ini cukup baru dan masih berada pada tahap eksperimental. Meskipun kami berusaha menyediakan API yang stabil, ada kemungkinan beberapa bagian kecil dari API publik akan berubah di masa depan.

Cara Kerja Proses: Gambaran Singkat

Cara Kerja Proses: Bekerja dengan Kode

Membuat instans model kosong

Alat pertama yang diperkenalkan oleh Accelerate untuk membantu menangani model besar adalah context manager init_empty_weights() yang membantu Anda menginisialisasi model tanpa menggunakan RAM, sehingga langkah 1 dapat dilakukan pada model berukuran berapa pun. Berikut cara kerjanya:

from accelerate import init_empty_weights

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

Sebagai contoh:

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

initializes sebuah model kosong dengan sedikit lebih dari 100 miliar parameter. Di balik layar, ini bergantung pada perangkat meta yang diperkenalkan dalam PyTorch 1.9. Selama inisialisasi di bawah context manager, setiap kali sebuah parameter dibuat, parameter tersebut langsung dipindahkan ke perangkat tersebut.

Anda tidak dapat memindahkan model yang diinisialisasi seperti ini ke CPU atau perangkat lain secara langsung, karena model tersebut tidak memiliki data apa pun. Selain itu, sangat mungkin bahwa forward pass dengan model kosong tersebut akan gagal, karena tidak semua operasi didukung pada meta device.

Checkpoint terbagi

Model Anda mungkin sangat besar sehingga bahkan satu salinannya pun tidak muat di RAM. Hal itu tidak berarti model tersebut tidak dapat dimuat: jika Anda memiliki satu atau beberapa GPU, ini berarti ada lebih banyak memori yang tersedia untuk menyimpan model Anda. Dalam kasus ini, lebih baik jika checkpoint Anda dipecah menjadi beberapa file yang lebih kecil yang kami sebut sebagai checkpoint shards.

Accelerate akan menangani checkpoint yang di-shard selama Anda mengikuti format berikut: checkpoint Anda harus berada dalam sebuah folder, dengan beberapa file yang berisi state dict parsial, dan harus ada index dalam format JSON yang berisi kamus yang memetakan nama parameter ke file yang berisi bobotnya. Anda dapat dengan mudah melakukan sharding model Anda dengan save_model(). Misalnya, kita bisa memiliki folder yang berisi:

first_state_dict.bin
index.json
second_state_dict.bin

dengan index.json sebagai file berikut:

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

dan first_state_dict.bin yang berisi bobot untuk "linear1.weight" dan "linear1.bias", second_state_dict.bin untuk "linear2.weight" dan "linear2.bias"

Memuat bobot

Alat kedua yang diperkenalkan oleh Accelerate adalah fungsi load_checkpoint_and_dispatch(), yang memungkinkan Anda memuat checkpoint ke dalam model kosong Anda. Fungsi ini mendukung checkpoint penuh (satu file yang berisi seluruh state dict) maupun checkpoint yang terbagi (sharded). Fungsi ini juga akan secara otomatis mendistribusikan bobot tersebut ke perangkat yang tersedia (GPU, RAM CPU), sehingga jika Anda memuat checkpoint yang terbagi, penggunaan RAM maksimum akan sama dengan ukuran shard terbesar.

Jika Anda ingin menggunakan inferensi model besar dengan model Transformers, lihat dokumentasi ini.

Berikut adalah cara kita dapat menggunakan ini untuk memuat model GPT2-1.5B.

Mari kita unduh versi terbagi (sharded) dari model ini.

pip install huggingface_hub

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

Untuk menginisialisasi model, kami akan menggunakan pustaka 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)

Lalu, muat checkpoint yang baru saja diunduh dengan:

from accelerate import load_checkpoint_and_dispatch

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

Dengan meneruskan device_map="auto", kita memberi tahu Accelerate untuk menentukan secara otomatis di mana setiap lapisan model harus ditempatkan berdasarkan sumber daya yang tersedia:

• pertama, kami menggunakan ruang maksimum yang tersedia pada GPU
• jika kami masih membutuhkan ruang, kami menyimpan sisa bobot di CPU
• jika RAM tidak mencukupi, sisa bobot disimpan di hard drive sebagai tensor memory-mapped

no_split_module_classes

Parameter ini akan menunjukkan bahwa beberapa modul dengan nama "Block" tidak boleh dipecah ke perangkat yang berbeda. Anda harus mengatur semua blok di sini yang mencakup koneksi residual dari jenis apa pun.

The device_map

Anda dapat melihat device_map yang dipilih oleh Accelerate dengan mengakses atribut hf_device_map dari model Anda:

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}

Sangat mungkin untuk membuat peta perangkat Anda sendiri untuk lapisan yang akan digunakan, dengan menentukan perangkat GPU yang akan dipakai (sebuah angka), "cpu", atau "disk" dan meneruskannya:

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
)

Jalankan model

Sekarang setelah kita melakukan ini, model kita tersebar di beberapa perangkat, dan mungkin juga di hard drive. Namun, model tersebut masih dapat digunakan sebagai model PyTorch biasa:

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())

Di balik layar, Accelerate menambahkan hook ke model, sehingga:

• pada setiap lapisan, input ditempatkan pada perangkat yang tepat (sehingga meskipun model Anda tersebar di beberapa GPU, tetap berfungsi)
• untuk bobot yang di-offload ke CPU, bobot tersebut dipindahkan ke GPU tepat sebelum forward pass dan dibersihkan tepat setelahnya
• untuk bobot yang di-offload ke hard drive, bobot tersebut dimuat ke RAM lalu dipindahkan ke GPU tepat sebelum forward pass dan dibersihkan tepat setelahnya

Dengan cara ini, model Anda dapat menjalankan inferensi meskipun tidak muat di salah satu GPU atau RAM CPU!

Ini hanya mendukung inferensi model Anda, bukan pelatihan. Sebagian besar komputasi terjadi di balik context manager torch.no_grad() untuk menghindari penggunaan memori GPU dengan aktivasi perantara.

Merancang peta perangkat

Anda dapat membiarkan Accelerate menangani perhitungan peta perangkat dengan mengatur device_map ke salah satu opsi yang didukung ("auto", "balanced", "balanced_low_0", "sequential") atau membuatnya sendiri jika Anda ingin lebih banyak kendali atas di mana setiap lapisan harus ditempatkan.

Anda dapat menurunkan semua ukuran model (dan dengan demikian menghitung device_map) pada model yang berada di meta device.

Semua opsi akan menghasilkan hasil yang sama ketika Anda tidak memiliki memori GPU yang cukup untuk menampung seluruh model (yaitu memuat sebanyak mungkin komponen ke GPU, lalu memindahkan bobot ke CPU atau bahkan ke disk jika RAM tidak mencukupi).

Ketika Anda memiliki memori GPU yang tersedia lebih besar daripada ukuran model, berikut adalah perbedaan antara masing-masing opsi:

• "auto" dan "balanced" membagi model secara merata di semua GPU yang tersedia, sehingga memungkinkan Anda menggunakan ukuran batch lebih dari 1.
• "balanced_low_0" membagi model secara merata di semua GPU kecuali yang pertama, dan hanya menempatkan di GPU 0 apa yang tidak muat di GPU lainnya. Opsi ini sangat berguna ketika Anda perlu menggunakan GPU 0 untuk memproses output, seperti saat menggunakan fungsi generate untuk model Transformers
• "sequential" akan mengisi GPU 0 sebanyak yang bisa, lalu beralih ke GPU 1 dan seterusnya (jadi tidak akan menggunakan GPU terakhir jika tidak diperlukan).

Opsi "auto" dan "balanced" saat ini menghasilkan hasil yang sama, tetapi perilaku "auto" mungkin berubah di masa depan jika kami menemukan strategi yang lebih masuk akal, sementara "balanced" akan tetap stabil.

Pertama, perhatikan bahwa Anda dapat membatasi memori yang digunakan pada setiap GPU dengan menggunakan argumen max_memory (tersedia di infer_auto_device_map() dan di semua fungsi yang menggunakannya). Saat mengatur max_memory, Anda harus meneruskan kamus yang berisi pengidentifikasi GPU (misalnya 0, 1, dst.) dan kunci "cpu" untuk RAM maksimum yang ingin Anda gunakan untuk CPU offload. Nilainya dapat berupa bilangan bulat (dalam byte) atau string yang mewakili angka dengan satuannya, seperti "10GiB" atau "10GB".

Berikut adalah contoh di mana kami tidak ingin menggunakan lebih dari 10GiB pada masing-masing dari dua GPU dan tidak lebih dari 30GiB RAM CPU untuk bobot model:

from accelerate import infer_auto_device_map

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

Saat alokasi pertama terjadi di PyTorch, kernel CUDA dimuat yang memakan sekitar 1-2GB memori tergantung pada GPU. Oleh karena itu, memori yang dapat digunakan selalu lebih sedikit daripada ukuran aktual GPU. Untuk melihat berapa banyak memori yang sebenarnya digunakan, lakukan torch.ones(1).cuda() dan periksa penggunaan memori.

Oleh karena itu, saat Anda membuat memory maps dengan max_memory, pastikan untuk menyesuaikan memori yang tersedia secara sesuai untuk menghindari error out-of-memory.

Selain itu, jika Anda melakukan beberapa operasi tambahan pada output Anda tanpa mengembalikannya ke CPU (misalnya di dalam metode generate dari Transformers) dan jika Anda menempatkan input Anda di GPU, GPU tersebut akan mengonsumsi lebih banyak memori daripada yang lain (Accelerate selalu mengembalikan output ke perangkat input). Oleh karena itu, jika Anda ingin mengoptimalkan ukuran batch maksimum dan Anda memiliki banyak GPU, berikan memori yang lebih sedikit ke GPU pertama. Misalnya, dengan BLOOM-176B pada setup 8x80 A100, peta yang mendekati ideal adalah:

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

seperti yang dapat Anda lihat, kami memberikan 7 GPU sisanya sekitar 50% lebih banyak memori daripada GPU 0.

Jika Anda memilih untuk merancang device_map sepenuhnya sendiri, ia harus berupa kamus dengan kunci berupa nama modul dari model Anda dan nilai berupa pengidentifikasi perangkat yang valid (misalnya bilangan bulat untuk GPU) atau "cpu" untuk offload CPU, "disk" untuk offload disk. Kunci-kunci tersebut harus mencakup seluruh model, dan Anda dapat mendefinisikan peta perangkat Anda sesuai keinginan: misalnya, jika model Anda memiliki dua blok (misalnya block1 dan block2) yang masing-masing berisi tiga lapisan linear (misalnya linear1, linear2 dan linear3), peta perangkat yang valid dapat berupa:

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

satu lagi yang valid bisa berupa:

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

Di sisi lain, yang ini tidak valid karena tidak mencakup setiap parameter model:

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

Untuk mencapai efisiensi tertinggi, pastikan peta perangkat Anda menempatkan parameter pada GPU secara berurutan (misalnya, jangan menempatkan salah satu bobot pertama di GPU 0, lalu bobot di GPU 1, dan bobot terakhir kembali ke GPU 0) untuk menghindari banyak transfer data antar GPU.

Hanya offload CPU

Jika Anda ingin memindahkan model ke CPU, Anda dapat menggunakan cpu_offload(). Akibatnya, semua parameter model akan dipindahkan dan hanya satu salinan state dict model yang akan disimpan. Selama forward pass, parameter akan diekstrak dari state dict tersebut, dipindahkan ke perangkat eksekusi, dan diteruskan sesuai kebutuhan, kemudian dipindahkan kembali.

cpu_offload(model, execution_device)

Anda juga dapat menggunakan cpu_offload_with_hook(). Fungsi ini akan memindahkan model ke CPU dan mengembalikannya ke perangkat eksekusi saat dieksekusi. Perbedaannya dengan cpu_offload() adalah model tetap berada di perangkat eksekusi setelah forward dan hanya dipindahkan kembali ketika metode offload dari hook yang dikembalikan dipanggil. Selain itu, cpu_offload_with_hook() lebih performan tetapi kurang hemat memori. Ini berguna untuk pipeline yang menjalankan model dalam loop:

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()

Hanya offload disk

Untuk melakukan disk offload, Anda dapat menggunakan disk_offload(). Akibatnya, semua parameter model akan di-offload sebagai array memory-mapped di folder yang ditentukan. Selama forward pass, parameter akan diakses dari folder tersebut dan dipindahkan ke perangkat eksekusi yang diberikan sesuai kebutuhan, kemudian di-offload kembali.

disk_offload(model, offload_dir, execution_device)

Batasan dan pengembangan lebih lanjut

Kami menyadari keterbatasan saat ini dalam API:

• infer_auto_device_map() (atau device_map="auto" di load_checkpoint_and_dispatch()) berusaha memaksimalkan RAM GPU dan CPU yang terlihat tersedia saat Anda mengeksekusinya. Meskipun PyTorch sangat baik dalam mengelola RAM GPU secara efisien (dan mengembalikannya saat tidak diperlukan), hal ini tidak sepenuhnya berlaku untuk Python dan RAM CPU. Oleh karena itu, peta perangkat yang dihitung secara otomatis mungkin terlalu membebani CPU. Pindahkan beberapa modul ke perangkat disk jika Anda mengalami crash akibat kekurangan RAM.
• infer_auto_device_map() (atau device_map="auto" di load_checkpoint_and_dispatch()) mengalokasikan perangkat secara berurutan (untuk menghindari perpindahan data bolak-balik), sehingga jika lapisan pertama Anda lebih besar dari ukuran GPU yang Anda miliki, semuanya akan berakhir di CPU/Disk.
• load_checkpoint_and_dispatch() dan load_checkpoint_in_model() saat ini tidak melakukan pemeriksaan apa pun terhadap kebenaran state dict Anda dibandingkan dengan model Anda (hal ini akan diperbaiki pada versi mendatang), sehingga Anda mungkin akan mendapatkan beberapa error aneh jika mencoba memuat checkpoint dengan kunci yang tidak cocok atau hilang.
• Paralelisme model yang digunakan saat model Anda dipecah ke beberapa GPU bersifat naif dan tidak dioptimalkan, yang berarti hanya satu GPU yang bekerja pada waktu tertentu sementara GPU lainnya menganggur.
• Ketika bobot di-offload ke CPU/disk keras, tidak ada pre-fetching (untuk saat ini, kami akan mengerjakan ini untuk versi mendatang) yang berarti bobot dipindahkan ke GPU hanya ketika dibutuhkan, bukan sebelumnya.
• Offloading ke hard drive mungkin sangat lambat jika perangkat keras yang Anda gunakan tidak memiliki komunikasi cepat antara disk dan CPU (seperti NVMes).