Sequence parallel di 🤗 accelerate

Panduan ini akan membahas dasar-dasar penggunaan sequence parallelism di 🤗accelerate.

Lihat juga Context Parallelism yang sangat terkait.

Mengapa paralelisme urutan?

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.

Ulysses Sequence parallelism memungkinkan kita untuk membagi (shard) input ke dalam komputasi attention di sepanjang dimensi sekuens dan menghitung attention secara normal, tetapi hanya menggunakan sebagian (slice) dari attention heads di setiap GPU. Dengan cara ini, kita dapat melatih model dengan sekuens panjang, dengan beberapa alat tambahan, hingga skala panjang sekuens 15M+. Untuk melihat cara melengkapi Ulysses SP dengan TiledMLP, Liger-Kernel, Activation checkpoint offload ke cpu, dan beberapa trik lainnya, silakan merujuk ke makalah: Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences.

Bagaimana Ulysses SP berbeda dari FSDP CP

Pada dokumen Context Parallelism, Anda dapat mempelajari cara menerapkan teknologi lain yang disebut Context Parallelism, yang juga memotong pada dimensi sekuens tetapi menggunakan Ring Attention alih-alih memotong pada dimensi head.

Artikel-artikel berikut memberikan penjelasan yang sangat rinci mengenai perbedaan antara kedua teknologi tersebut:

• https://insujang.github.io/2024-01-11/tensor-parallelism-and-sequence-parallelism-detailed-analysis/
• https://huggingface.co/blog/exploding-gradients/ulysses-ring-attention

Ringkasan singkat yang diadaptasi dari salah satu artikel:

• Ulysses SP memiliki overhead komunikasi yang relatif rendah, tetapi dibatasi oleh jumlah Attention Heads, sehingga memiliki persyaratan tertentu untuk topologi jaringan (jumlah attention heads harus dapat dibagi habis oleh jumlah GPU yang berpartisipasi untuk satu replika). Komunikasi all-to-all sangat sensitif terhadap latensi dan memerlukan Deepspeed.
• Komunikasi ring P2P dari FSDP CP Ring-Attention tidak memiliki persyaratan keterbagian yang disebutkan sebelumnya, tetapi memiliki volume komunikasi yang lebih tinggi.

Akhirnya, seharusnya dimungkinkan untuk menggabungkan SP + CP seperti yang dijelaskan dalam makalah USP: A Unified Sequence Parallelism Approach for Long Context Generative AI untuk mendukung panjang sekuens yang bahkan lebih panjang, meskipun hal ini belum terintegrasi ke dalam 🤗accelerate.

Backend paralelisme urutan yang didukung

Saat ini, satu-satunya backend paralelisasi urutan adalah deepspeed, yang berasal dari Ulysses SP yang dimodernisasi dan merupakan bagian dari teknologi Arctic Long Sequence Training. Tersedia juga tutorial jika Anda ingin mengintegrasikannya langsung ke kode Anda sendiri.

Bagaimana cara menggunakan paralelisme urutan?

from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

+# Example: 4 GPUs with sp_size=4, dp_shard_size=1
+# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
+     sp_backend="deepspeed",
+     sp_size=4,
+     dp_shard_size=1,  # Explicit: no data parallelism
+     sp_handler=DeepSpeedSequenceParallelConfig(
+         sp_seq_length_is_variable: true,
+         sp_attn_implementation="sdpa",
+     ),
+ )

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

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

accelerate launch --parallelism-config-sp-size 8  ...

Anda juga dapat mengatur sp_size dan konfigurasi lainnya melalui perintah accelerate config, yang akan menyimpannya ke dalam file konfigurasi accelerate Anda, sehingga Anda tidak perlu meneruskannya setiap kali menjalankan skrip.

sequence parallelism dikombinasikan dengan data parallelism. Tidak memerlukan GPU tambahan. Jadi jika Anda memiliki 8 GPU, Anda dapat melakukan: --parallelism-config-dp-shard-size 8 --parallelism-config-sp-size 8. Atau Anda dapat menggunakan kelas ParallelismConfig untuk menyetelnya secara programatik.

Penting: Anda harus memastikan dp_replicate_size × dp_shard_size × sp_size = num_processes. Misalnya, dengan 8 GPU dan sp_size=8, Anda membutuhkan dp_shard_size=1 (karena 1 × 1 × 8 = 8). Dengan 4 GPU dan sp_size=2, Anda dapat menggunakan dp_shard_size=2 (karena 1 × 2 × 2 = 4) untuk paralelisme 2D.

Konfigurasi backend ALST/Ulysses SP

ALST/UlyssesSP menerapkan paralelisme sekuens menggunakan paralelisme kepala perhatian, seperti dijelaskan dalam makalah ini. Untuk kesederhanaan, kami menggunakan kembali konsep dan pengaturan paralelisme sekuens, yang, dari sudut pandang pengguna, sama: beberapa GPU digunakan untuk memproses satu batch.

Untuk memberikan gambaran tentang apa yang ALST memungkinkan - fitur ini memungkinkan kami melatih dalam bf16 dengan 500K token pada satu GPU H100, 3,7M pada satu node, dan 15M pada Llama-8B menggunakan hanya empat node. Fitur dari HF Accelerate ini hanya mengaktifkan 1 dari 3 komponen ALST, sehingga panjang sekuens yang dapat dicapai akan lebih kecil. Anda perlu mengaktifkan TiledMLP, Activation checkpoint offload to CPU, dan beberapa hal lain untuk mendapatkan kekuatan penuh ALST. Untuk detailnya, silakan merujuk pada tutorial ini.

Untuk mengonfigurasi backend deepspeed:

# Example: 4 GPUs with sp_size=4, dp_shard_size=1
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_shard_size=1,  # Explicit: no data parallelism
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length=256,
        sp_seq_length_is_variable=True,
        sp_attn_implementation="sdpa",
    ),
)
accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

• sp_backend: diset ke deepspeed di sini
• sp_size adalah derajat paralelisasi sekuens - dalam contoh di atas nilainya 4, sehingga 4 GPU akan digunakan untuk memproses satu batch (sambil melakukan DP=4 di atas GPU yang sama)
• sp_seq_length dan sp_seq_length_is_variable digunakan untuk menangani panjang sekuens. Jika sp_seq_length_is_variable=True, backend akan bekerja dengan panjang sekuens yang dapat berubah antar batch, di mana nilai sp_seq_length dapat diatur ke apa pun yang habis dibagi derajat paralel sekuens atau tidak diatur sama sekali. Dalam kasus ini, pada setiap forward, variabel sekuens akan diturunkan dari input. Jika False, maka seq_length harus sesuai dengan dimensi panjang sekuens batch, yang kemudian harus di-pad agar selalu sama. Nilai defaultnya adalah True.
• sp_attn_implementation adalah salah satu dari sdpa, flash_attention_2, atau flash_attention_3. Implementasi paralel urutan ini menggunakan position_ids alih-alih attention_mask, sehingga eager tidak dapat berfungsi di sini sampai mendukung penggunaan position_ids. Selain itu, perhatikan bahwa sdpa tidak menangani beberapa sampel yang digabungkan menjadi satu dengan benar; ia akan memperhatikan seluruh sampel sebagai satu kesatuan. Jika sampel tidak digabungkan, sdpa akan berfungsi dengan benar. Oleh karena itu, Flash Attention seharusnya menjadi pilihan ideal karena selalu berfungsi.

Alih-alih mengatur nilai-nilai ini dalam objek DeepSpeedSequenceParallelConfig, Anda juga dapat menggunakan variabel lingkungan untuk mencapai hal yang sama - di sini mereka sesuai dengan akhir dari daftar di atas.

• PARALLELISM_CONFIG_SP_BACKEND
• PARALLELISM_CONFIG_SP_SEQ_LENGTH
• PARALLELISM_CONFIG_SP_SEQ_LENGTH_IS_VARIABLE
• PARALLELISM_CONFIG_SP_ATTN_IMPLEMENTATION

Jika tidak diteruskan dalam kode, sp_size dapat diatur melalui argumen CLI --parallelism_config_sp_size. Hal yang sama berlaku untuk argumen lainnya. Anda juga dapat menggunakan gaya konfigurasi file accelerate config, misalnya, untuk 2 GPU:

distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 2
parallelism_config:
  parallelism_config_dp_replicate_size: 1
  parallelism_config_dp_shard_size: 1  # Must satisfy: 1 × 1 × 2 = 2 num_processes
  parallelism_config_sp_size: 2
  parallelism_config_sp_backend: deepspeed
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: sdpa

Seperti yang disebutkan sebelumnya, paralelisasi sekuens Ulysses biasanya ditumpangkan dengan paralelisasi data - rank yang sama digunakan untuk memberi umpan aliran data unik dan juga melakukan Ulysses Sequence Parallelism. Namun, Anda juga dapat membuat replika seperti ini:

# Example: 4 GPUs with 2D parallelism (SP=2, DP=2)
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 2 × 1 × 2 = 4 GPUs
parallelism_config = ParallelismConfig(
    dp_replicate_size=2,
    dp_shard_size=1,  # Explicit: no sharding within replicas
    sp_size=2,
    sp_backend="deepspeed",
    sp_handler=DeepSpeedSequenceParallelConfig(...),
)

Di sini kami menggunakan 4 gpus, dengan 2 replika sequence parallelism. Deepspeed-ZeRO adalah yang menggerakkan data parallelism di sini.

Perlu dicatat bahwa banyak hal tersembunyi di dalam UlyssesSPDataLoaderAdapter. Komponen ini bekerja di balik layar, membungkus objek DataLoader asli Anda, namun Anda perlu menyadarinya jika menemui masalah. Ia juga secara otomatis menyuntikkan shift_labels yang benar ke dalam kamus batch, sebelum batch dipecah di antara rank yang berpartisipasi.

Sekarang, satu-satunya bagian tersisa untuk mulai menggunakan ALST/UlyssesSP adalah mengagregasi loss di seluruh rank menggunakan all_gather yang dapat dibedakan agar grads benar. Kode berikut melakukannya, sambil juga mengecualikan token yang di-mask dengan -100, untuk mendapatkan rata-rata yang benar:

sp_size = parallelism_config.sp_size if parallelism_config is not None else 1
if sp_size > 1:
    sp_group = accelerator.torch_device_mesh["sp"].get_group()
    sp_world_size = parallelism_config.sp_size

# Normal training loop
for iter, batch in enumerate(dl):
    optimizer.zero_grad()

    batch = move_to_device(batch, model.device)

    # The model automatically receives shift_labels via **kwargs and uses it for loss computation.
    # Both standard transformers models and Liger-patched models handle this correctly.
    outputs = model(**batch)
    loss = outputs.loss
    shift_labels = batch["shift_labels"]

    if sp_size > 1:
        # differentiable weighted per-shard-loss aggregation across ranks
        losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
        # special dealing with SFT that has prompt tokens that aren't used in loss computation
        good_tokens = (shift_labels != -100).view(-1).sum()
        good_tokens_per_rank = torch.distributed.nn.functional.all_gather(
            good_tokens, group=sp_group
        )
        # Skip ranks with zero valid tokens to avoid NaN contamination (NaN * 0 = NaN)
        total_loss = sum(
            losses_per_rank[rank] * good_tokens_per_rank[rank]
            for rank in range(sp_world_size)
            if good_tokens_per_rank[rank] > 0
        )
        total_good_tokens = sum(good_tokens_per_rank)
        loss = total_loss / max(total_good_tokens, 1)

    if rank == 0: accelerator.print(f"{iter}: {loss=}")
    accelerator.log(dict(train_loss=loss, step=iter))

    accelerator.backward(loss)
    optimizer.step()

Perhatikan bahwa model secara otomatis menangani shift_labels ketika ada dalam batch. Forward pass model menerima shift_labels melalui **kwargs dan meneruskannya ke fungsi loss, yang kemudian menghitung loss dengan benar untuk sequence parallelism. Jika Anda menggunakan Liger Kernel, ia juga menangani shift_labels secara mulus dan menghitung loss dengan cara yang sangat hemat memori. Liger sangat direkomendasikan untuk panjang sequence yang panjang, karena membebaskan memori GPU dengan menggunakan operasi fused (misalnya, perhitungan fused logit-loss yang tidak pernah mematerialisasi tensor logits penuh di memori).

Jika Anda ingin melihat apa yang dilakukan HF Accelerate di balik layar, silakan baca tutorial integrasi lengkap ini.

Untuk contoh loop pelatihan Accelerate dengan ALST/UlyssesSP yang diaktifkan, lihat examples/alst_ulysses_sequence_parallelism.

[!Peringatan]

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

Karena ini adalah backend Deepspeed, konfigurasi Deepspeed yang biasa berlaku, sehingga Anda juga dapat menggabungkan sequence parallelism dengan offloading optimizer states dan/atau weights untuk membebaskan lebih banyak memori gpu dan memungkinkan panjang sequence yang bahkan lebih panjang. Teknologi ini telah diuji dan terbukti berfungsi dengan DeepSpeed ZeRO stage 2 dan 3.