Context Parallel di 🤗 accelerate

Panduan ini akan membahas dasar-dasar penggunaan context parallelism di 🤗accelerate. Bagi pembaca yang lebih penasaran, kami juga akan membahas beberapa aspek teknis di bagian-bagian selanjutnya.

Lihat juga Panduan Paralelisme Urutan yang sangat terkait.

Mengapa paralelisme konteks?

Dengan kemunculan model bahasa besar, dan baru-baru ini model penalaran, panjang sekuens telah tumbuh dengan cepat. Hal ini, dikombinasikan dengan kompleksitas memori kuadratik dari attention, telah menimbulkan kebutuhan akan cara yang lebih efisien untuk melatih model dengan sekuens panjang. Dengan panjang sekuens 128k, kebutuhan memori untuk matriks attention adalah 128k * 128k * 2 bytes * num_heads = ~32 GB * num_heads untuk presisi bf16, dengan asumsi implementasi attention vanilla. Memang, dengan penggunaan flash attention atau SDPA yang tidak mematerialkan bobot attention tersebut, kebutuhan ini berkurang secara drastis, tetapi pertumbuhan kebutuhan memori tetap cukup signifikan.

Context parallelism memungkinkan kita untuk membagi (shard) input ke dalam komputasi attention di sepanjang dimensi sekuens dan menghitung attention secara paralel di beberapa GPU. Dengan ini, kita dapat melatih model dengan sekuens panjang, yang berpotensi berskala hingga panjang sekuens 1 juta+.

Bagaimana cara menggunakan paralelisme konteks?

from accelerate.utils import ParallelismConfig, TorchContextParallelConfig

+ cp_config = TorchContextParallelConfig(
+       cp_comm_strategy="alltoall", # no need to use cp_config at all, if you want to use the default "allgather"
+ )

+ parallelism_config = ParallelismConfig(
+     cp_size=8,
+     cp_handler=cp_config,  # or just cp_size=8, if you want to use the default "allgather"
+ )

accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

Seperti fitur lain di 🤗accelerate, Anda juga dapat mengaktifkan paralelisme konteks dengan meneruskan flag yang sesuai ke accelerate launch. Dalam hal ini, tidak ada perbedaan:

accelerate launch --parallelism-config-cp-size 8 --parallelism-config-cp-comm-strategy [allgather|alltoall] ...

Anda juga dapat mengatur cp_size dan cp_comm_strategy dalam perintah accelerate config, yang akan menyimpannya di file konfigurasi accelerate Anda, sehingga Anda tidak perlu meneruskannya setiap kali menjalankan skrip Anda.

Context parallelism kompatibel dengan strategi paralelisasi lainnya, seperti data parallelism, tensor parallelism, dan FSDP2. Anda dapat menggabungkannya dengan mudah dengan mengatur ukuran paralelisasi ke nilai yang diinginkan, misalnya --parallelism-config-dp-size 8 --parallelism-config-tp-size 2 --parallelism-config-cp-size 8. Atau Anda dapat menggunakan kelas ParallelismConfig untuk mengaturnya secara programatik.

Context parallelism sangat terikat dengan FSDP2, yang dapat Anda pelajari lebih lanjut dalam pengenalan FSDP2. Artinya, context parallelism hanya berfungsi jika Anda menggunakan FullyShardedDataParallelPlugin atau --use-fsdp dengan version yang diatur ke 2 pada program Anda. Jika tidak ada FSDP2 yang digunakan, error akan dilempar.

Context parallelism hanya berfungsi dengan SDPA dan hanya tanpa mask atau causal mask. Kami tidak dapat mendeteksi hal ini dengan benar untuk Anda, sehingga menjadi tanggung jawab Anda untuk memastikan bahwa Anda menggunakan SDPA tanpa mask atau causal mask. Jika Anda menggunakan implementasi attention lainnya, akan muncul error.

Setelah mengaktifkan context parallelism dengan metode yang disebutkan di atas, Anda dapat menerapkannya ke dalam training loop Anda. Kami menyediakan wrapper tipis di sekitar torch.distributed.tensor.experimental.context_parallel yang dapat Anda gunakan dalam training loop, yang mengabstraksi sebagian kompleksitas penggunaannya (lebih lanjut tentang ini nanti). Untuk meminimalkan perubahan yang harus Anda lakukan pada training loop, kami menyediakan context manager yang bersifat noop jika context parallelism tidak diaktifkan, dan menerapkan context parallelism jika diaktifkan. Dengan cara ini, Anda dapat menggunakannya dalam training loop tanpa mengubah kode apa pun berdasarkan konfigurasi parallelism Anda. Anda dapat menggunakannya sebagai berikut:

for batch in dataloader:
    with accelerator.maybe_context_parallel(
        buffers=[batch["input_ids"], batch["attention_mask"]],
        buffer_seq_dims=[1, 1],
        no_restore_buffers={batch["input_ids"], batch["labels"]},
    ):
        outputs = model(**batch)
        ...

Context manager ini harus dibuat ulang pada setiap langkah pelatihan, seperti yang ditunjukkan pada contoh di atas. Hal ini sangat penting untuk dilakukan.

Hal ini dapat memperluas ukuran konteks Anda hingga panjang sekuens 1M+. Di bawah ini, kami menampilkan kecepatan dan penggunaan memori dari paralelisme konteks untuk ukuran konteks hingga 256k. Dapat dilihat bahwa ketika kami menggandakan ukuran konteks dan jumlah GPU, kami dapat mencapai penggunaan memori yang konsisten, yang berpotensi memungkinkan penskalaan panjang konteks tanpa batas.

Gambar 1: Penggunaan memori dan kecepatan paralelisme konteks untuk ukuran konteks hingga 256k.

Contoh-contoh ini dibuat dengan skrip yang dapat Anda temukan di folder examples. Untuk menjalankan contoh pada 8 GPU H100 (panjang sekuens 128k), Anda dapat menggunakan perintah berikut:

Disalin

accelerate launch --use-fsdp --fsdp-activation-checkpointing=TRUE examples/torch_native_parallelism/nd_parallel.py --cp-size=8 --sequence-length=128000

Antarmuka Accelerate

context manager menerima beberapa argumen yang digunakan untuk mengonfigurasi context parallelism.

• buffers: Ini adalah daftar tensor yang akan dipecah (sharded) di sepanjang dimensi urutan. Tensor-tensor ini biasanya berupa input ids, labels, dan attention mask.
• buffer_seq_dims: Ini adalah daftar bilangan bulat yang menentukan dimensi sekuens dari buffer, sesuai dengan urutan daftar buffers. Jika Anda meneruskan buffers=[input_ids, shift_labels] dengan keduanya memiliki bentuk [batch_size, sequence_length], Anda akan meneruskan buffer_seq_dims=[1, 1] karena dimensi sekuens adalah dimensi kedua dari tensor. Hal ini diperlukan untuk perhitungan output model yang benar.
• no_restore_buffers: Implementasi context parallelism memodifikasi buffer secara in-place, mengubahnya menjadi torch.distributed.tensor.Dtensor. Setelah context manager keluar, kernel komunikasi perlu diluncurkan untuk mengembalikan buffer ke keadaan aslinya (biasanya all-gather). Proses ini memakan waktu, sehingga disarankan untuk meneruskan tensor yang sama seperti pada argumen buffers, untuk menghindari komunikasi yang tidak perlu, kecuali Anda yakin bahwa Anda perlu menggunakan buffer setelah context manager keluar.

Context parallelism tidak kompatibel dengan labels yang merupakan salinan dari input_ids, yang dapat digeser oleh model dari 🤗 transformers untuk memungkinkan causal language modeling secara mandiri. Bayangkan kasus ini: labels = [l1, l2, l3, l4, … li] jika kita menerapkan context parallelism, setiap rank akan berakhir dengan sebagian dari labels, seperti ini: labels_rank_0 = [l1, l2], labels_rank_1 = [l3, l4], … setelah kode pemodelan transformers menggeser labels, hasilnya akan menjadi: labels_rank_0 = [l2, PAD], labels_rank_1 = [l3, PAD], … di mana PAD adalah token padding. Hal ini akan menghasilkan perhitungan loss yang salah, karena labels tidak lagi selaras dengan input. Karena alasan ini, Anda perlu menggeser labels secara manual sebelum meneruskannya ke model

Opsi yang dapat dikonfigurasi

Accelerate hanya menyediakan satu opsi untuk mengonfigurasi paralelisme konteks (kecuali untuk cp_size)

• cp_comm_strategy: Metode rotasi yang digunakan untuk shard. Kami sangat merekomendasikan untuk tetap menggunakan "allgather", karena kemungkinan besar akan lebih unggul daripada "alltoall" dalam kebanyakan kasus.

Ukuran paralel konteks cukup jelas dari namanya, yaitu jumlah rank di mana input akan di-shard. Rotasi shard paralel konteks mendefinisikan bagaimana shard dari input dirotasikan di antara rank. Kami akan membahas 2 opsi ini secara lebih detail di bagian berikutnya.

Anda dapat melihat contoh end-to-end di file ND parallel example, di mana Anda dapat melatih model 8B dengan panjang konteks hingga 128k pada satu node 8xH100. Dengan menggunakan pelatihan multi-node, Anda dapat menskalakan ini ke panjang sekuens 1M+ pada beberapa GPU. Anda juga dapat menggabungkannya secara mulus dengan strategi paralelisasi lainnya sesuai kebutuhan Anda.

Detail teknis

Bagian ini cukup teknis, jadi jika Anda tidak perlu mempelajari bagian internal dari paralelisme konteks, Anda dapat melewatinya dan langsung mulai membangun 🚀

Kami akan menggunakan kata shard secara luas di bagian-bagian berikut, jadi mari kita definisikan terlebih dahulu. Jika kita menyebut tensor sharded di sepanjang dimensi Dth, di sepanjang N rank, yang dimaksud adalah tensor ini dipecah menjadi N bagian, di mana setiap bagian tensor memiliki bentuk [..., D//N, ...].

Jadi, bagaimana cara kerjanya?

Context parallelism works on sharding the Q, K and V matrices across the sequence dimension. Each rank has its assigned shard of Q, let’s call it Q_i. This matrix stays only on this rank, during the whole computation. Similarly, each rank has its own shard of K and V, let’s call them K_i and V_i. Then, each rank calculates attention with its own shard of Q_i, K_i and V_i, let’s call it attn_i. During this computation, a communication kernel is launched to gather the Ks and Vs from all other ranks. What communication primitive is used, depends on the context_parallel_shard_rotation option. This way, each rank gets to calculate local attention, first with Q_i, K_i and V_i, then with K_j and V_j from all other ranks. As each rank holds Q, K and V matrices that are sharded across the sequence dimension, the resulting matrices are smaller and can fit on a single GPU.

Kita dapat memformalkan hal ini dalam pseudocode berikut:

comm_kernel = {"allgather": allgather, "alltoall": alltoall}[context_parallel_shard_rotation]
Qi, Ki, Vi = shard(Q, K, V, seq_dim)
attn[i] = attn(Qi, Ki, Vi)
for j in range(context_parallel_size):
    Kj, Vj = comm_kernel()
    attn[j] = attn(Qi, Kj, Vj) # [batch, num_heads, seq_len // context_parallel_size, head_dim]

final_attn = combine(attn)

all-to-all vs all-gather

all-gather

Jadi, apa perbedaan antara all-to-all dan all-gather? Pada all-gather, komunikasinya sangat sederhana. Setelah (atau lebih tepatnya, sebelum, karena biasanya prosesnya memakan waktu lebih lama) kami menghitung perhatian lokal attn_i, kami meluncurkan all-gather untuk mengumpulkan semua Ks dan Vs lainnya dari semua rank lain. Setelah komunikasi ini selesai, setiap rank memiliki semua Ks dan Vs dari rank lain, dan dapat menghitung perhatian dengan mereka secara berurutan. Dalam skenario ideal, all-gather selesai tepat pada saat perhitungan attn_i selesai. Namun, hal ini tidak pernah terjadi dalam praktiknya, sehingga overlap nyata yang ideal tercapai ketika seluruh attn_i di-overlap dengan sebagian dari komunikasi, kemudian untuk memulai perhitungan dengan K_j dan V_j, kami menunggu hingga all-gather selesai.

all-to-all

All-to-all, atau kadang disebut ring-rotation, memanfaatkan pola komunikasi berbentuk cincin. Setelah menyelesaikan komputasi attn_i, sebuah all-to-all diluncurkan untuk mengirim K_i dan V_i ke rank tetangga. Kami kemudian mengulang ini sebanyak context_parallel_size-1 kali, sehingga setiap rank melihat semua shard dari K dan V dari semua rank lain sekali. Dalam skenario ideal, kami mengambil prefetch shard K_i+1 dan V_i+1 dari rank tetangga dan komunikasi ini tumpang tindih persis dengan komputasi attn_i kami saat ini. Sekali lagi, secara realistis, tumpang tindih sempurna ini tidak pernah terjadi. Mengingat sifat pendekatan ini, jika kami tidak mencapai tumpang tindih sempurna, penaltinya jauh lebih besar daripada dengan all-gather.

Bagaimana memilih metode rotasi yang tepat?

Secara teori, all-to-all seharusnya menjadi pilihan yang lebih baik. Namun, dalam praktiknya, hal itu jarang terjadi. Oleh karena itu, kami menggunakan all-gather secara default, karena lebih mungkin mencapai kinerja yang lebih baik. Benchmark yang luas dari tim torchtitan juga menunjukkan bahwa all-to-all jarang melampaui kinerja all-gather. Meskipun demikian, kami tetap menyediakan kedua opsi, karena Anda mungkin menemukan salah satunya lebih baik untuk kasus penggunaan Anda.

Anda dapat melihat masalah ini secara langsung pada keluaran profiler di gambar di bawah:

Gambar 1: Bagian berwarna merah menunjukkan waktu idle saat kita menunggu kernel all-to-all selesai. Pada bar biru pertama yang disorot, terlihat bahwa prosesnya membutuhkan waktu ~250us untuk selesai, yang diulang N-1 kali untuk setiap pemanggilan attention, di mana N adalah ukuran context parallel.

Mengapa hanya FSDP2?

Kami hanya mendukung paralelisme konteks dengan FSDP2, karena kami membuat mesh gabungan dari context_parallel_size dan dp_shard_size untuk memanfaatkan potensi penuhnya. Cara kerjanya adalah: kami memecah model di seluruh mesh gabungan berukuran cp_size*dp_shard_size, yang memaksimalkan penghematan memori. Ini semacam “free lunch”, karena komunikasi FSDP sepenuhnya tumpang tindih dengan komputasi attention, seperti yang ditunjukkan pada gambar di bawah.

Gambar 2: Pada persegi panjang biru (Stream 23), Anda dapat melihat bahwa pre-fetch shard `FSDP` sepenuhnya tumpang tindih dengan komputasi attention (Stream 7), sedangkan pada persegi panjang merah (Stream 24), Anda dapat melihat bahwa kernel all-gather menghasilkan gelembung waktu idle, di mana stream komputasi kami (7) tidak aktif.

Pada gambar di atas, Anda juga dapat memperhatikan perbedaan antara all-to-all dan all-gather. Sementara pada all-to-all (Gambar 1), kami meluncurkan kernel komunikasi N-1 kali untuk setiap pemanggilan attention, pada all-gather (Gambar 2), kami hanya meluncurkan kernel komunikasi sekali. Hal ini menghasilkan gelembung yang lebih besar, tetapi hanya terjadi sekali per pemanggilan attention, sedangkan pada all-to-all, hal itu terjadi N-1 kali.

Penyaluran data dalam mesh gabungan

Kami memastikan untuk mendistribusikan batch data yang sama ke seluruh subgrup cp, sehingga hasilnya benar. (Artinya, setiap rank dalam subgrup cp menerima batch data yang sama.) Namun, kami juga mendistribusikan batch yang berbeda ke setiap rank dalam grup dp_shard. Bayangkan seperti ini:

# 8 GPUS, --dp_shard_size 4, --cp_size 2
# mesh = [[0, 1], [2, 3], [4, 5], [6, 7]]
# model is sharded across the whole mesh (each GPU holds 1/8 of the model)
# GPUs 0,1 = batch 0
# GPUs 2,3 = batch 1
... and so on.