Song song chuỗi trong 🤗 accelerate

Hướng dẫn này sẽ bao quát các kiến thức cơ bản về việc sử dụng song song chuỗi (sequence parallelism) trong 🤗accelerate.

Xem thêm Context Parallelism rất liên quan.

Tại sao cần song song hóa chuỗi?

Với sự ra đời của các mô hình ngôn ngữ lớn, và gần đây là các mô hình suy luận, độ dài chuỗi đã tăng lên nhanh chóng. Điều này, kết hợp với độ phức tạp bộ nhớ bậc hai của cơ chế chú ý, đã dẫn đến nhu cầu về các phương pháp hiệu quả hơn để huấn luyện các mô hình với chuỗi dài. Với độ dài chuỗi là 128k, yêu cầu bộ nhớ của ma trận chú ý là 128k * 128k * 2 bytes * num_heads = ~32 GB * num_heads cho độ chính xác bf16, với giả định về việc triển khai chú ý thông thường. Mặc dù, với việc sử dụng flash attention hoặc SDPA, vốn không tạo ra các trọng số chú ý này, con số này giảm đi đáng kể, nhưng sự gia tăng trong yêu cầu bộ nhớ vẫn là đáng kể.

Ulysses Sequence parallelism cho phép chúng ta chia nhỏ (shard) các đầu vào của phép tính attention dọc theo chiều sequence và tính toán attention bình thường, nhưng chỉ sử dụng một phần (slice) các attention heads trên mỗi GPU. Nhờ đó, chúng ta có thể huấn luyện các mô hình với sequence dài, cùng với một vài công cụ bổ sung, mở rộng lên độ dài sequence hơn 15M. Để xem cách bổ sung TiledMLP, Liger-Kernel, Activation checkpoint offload to cpu và một vài mẹo khác vào Ulysses SP, vui lòng tham khảo bài báo: Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences.

Ulysses SP khác với FSDP CP như thế nào

Trong tài liệu Context Parallelism, bạn có thể tìm hiểu về việc triển khai một công nghệ khác có tên là Context Parallelism, công nghệ này cũng chia nhỏ theo chiều dài chuỗi nhưng sử dụng Ring Attention thay vì chia nhỏ theo chiều đầu (head).

Các bài viết sau đây đi vào giải thích rất chi tiết về sự khác biệt giữa hai công nghệ:

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

Tóm tắt nhanh được điều chỉnh từ một trong các bài viết:

• Ulysses SP có chi phí truyền thông tương đối thấp, nhưng bị giới hạn bởi số lượng Attention Heads và do đó có một số yêu cầu nhất định về topology mạng (số lượng attention heads phải chia hết cho số lượng GPU tham gia của một replica duy nhất). Truyền thông all-to-all nhạy cảm với độ trễ và yêu cầu Deepspeed.
• Giao tiếp vòng P2P của FSDP CP Ring-Attention không có các yêu cầu chia hết như đã nêu ở trên, nhưng có lượng giao tiếp lớn hơn.

Cuối cùng, cần phải có khả năng kết hợp SP + CP như đã mô tả trong bài báo USP: A Unified Sequence Parallelism Approach for Long Context Generative AI để hỗ trợ độ dài chuỗi thậm chí dài hơn, mặc dù điều này chưa được tích hợp vào 🤗accelerate.

Các backend song song chuỗi được hỗ trợ

Hiện tại, backend duy nhất cho song song chuỗi là deepspeed, được phát triển từ Ulysses SP hiện đại hóa, một phần của công nghệ Arctic Long Sequence Training. Ngoài ra, có một hướng dẫn nếu bạn muốn tích hợp trực tiếp vào mã của mình.

Cách sử dụng song song hóa chuỗi?

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

Tương tự như mọi tính năng khác trong 🤗accelerate, bạn cũng có thể bật song song chuỗi bằng cách truyền các cờ tương ứng vào accelerate launch. Trong trường hợp này, không có gì khác biệt:

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

Bạn cũng có thể đặt sp_size và các cấu hình khác trong lệnh accelerate config, chúng sẽ được lưu trong tệp cấu hình accelerate của bạn, giúp bạn không cần truyền chúng mỗi lần khởi chạy script.

sequence parallelism kết hợp với data parallelism. Nó không yêu cầu thêm GPU. Vì vậy, nếu bạn có 8 GPU, bạn có thể thực hiện: --parallelism-config-dp-shard-size 8 --parallelism-config-sp-size 8. Hoặc bạn có thể sử dụng lớp ParallelismConfig để thiết lập chúng theo chương trình.

Quan trọng: Bạn phải đảm bảo dp_replicate_size × dp_shard_size × sp_size = num_processes. Ví dụ, với 8 GPU và sp_size=8, bạn cần dp_shard_size=1 (vì 1 × 1 × 8 = 8). Với 4 GPU và sp_size=2, bạn có thể dùng dp_shard_size=2 (vì 1 × 2 × 2 = 4) cho song song 2D.

Cấu hình backend ALST/Ulysses SP

ALST/UlyssesSP triển khai song song hóa chuỗi bằng cách sử dụng song song hóa đầu chú ý, như được giải thích trong bài báo này. Để đơn giản hóa, chúng tôi tái sử dụng khái niệm và cấu hình của song song hóa chuỗi, mà từ góc nhìn của người dùng thì giống hệt nhau: nhiều GPU được sử dụng để xử lý một lô (batch) duy nhất.

Để có cái nhìn về những gì ALST đã cho phép - nó cho phép chúng tôi huấn luyện ở bf16 với 500K token trên một GPU H100 đơn lẻ, 3,7M trên một node đơn lẻ, và 15M trên Llama-8B chỉ bằng bốn node. Tính năng của HF Accelerate này chỉ kích hoạt 1 trong 3 thành phần của ALST, do đó độ dài chuỗi có thể đạt được sẽ nhỏ hơn. Bạn sẽ cần bật TiledMLP, Activation checkpoint offload to CPU và một số thứ khác để tận dụng toàn bộ sức mạnh của ALST. Để biết chi tiết, vui lòng tham khảo tutorial này.

Để cấu hình 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: đặt thành deepspeed tại đây
• sp_size là độ song song của chuỗi - trong ví dụ trên, giá trị này là 4, do đó 4 GPU sẽ được sử dụng để xử lý một batch duy nhất (đồng thời thực hiện DP=4 trên cùng các GPU đó)
• sp_seq_length và sp_seq_length_is_variable được dùng để xử lý độ dài chuỗi. Nếu sp_seq_length_is_variable=True, backend sẽ hoạt động với độ dài chuỗi có thể thay đổi giữa các batch, trong trường hợp đó giá trị của sp_seq_length có thể được đặt thành bất kỳ giá trị nào chia hết cho bậc song song chuỗi hoặc không được đặt. Trong trường hợp này, ở mỗi lần forward, các biến chuỗi sẽ được suy ra từ đầu vào. Nếu False, thì seq_length phải khớp với chiều độ dài chuỗi của batch, và chiều này sẽ phải được đệm (padding) để luôn giữ nguyên. Giá trị mặc định là True.
• sp_attn_implementation là một trong sdpa, flash_attention_2 hoặc flash_attention_3. Bản triển khai song song chuỗi này sử dụng position_ids thay vì attention_mask, do đó eager không thể hoạt động ở đây cho đến khi nó hỗ trợ làm việc với position_ids. Ngoài ra, vui lòng lưu ý rằng sdpa không xử lý chính xác nhiều mẫu được kết hợp thành một; nó sẽ chú ý đến toàn bộ mẫu như một. Nếu các mẫu không được kết hợp, sdpa sẽ hoạt động chính xác. Do đó, Flash Attention nên là lựa chọn lý tưởng vì nó luôn hoạt động.

Thay vì đặt các giá trị này trong đối tượng DeepSpeedSequenceParallelConfig, bạn cũng có thể sử dụng các biến môi trường để đạt được kết quả tương tự - dưới đây là các biến tương ứng với cuối danh sách ở trên.

• PARALLELISM_CONFIG_SP_BACKEND
• PARALLELISM_CONFIG_SP_SEQ_LENGTH
• PARALLELISM_CONFIG_SP_SEQ_LENGTH_IS_VARIABLE
• PARALLELISM_CONFIG_SP_ATTN_IMPLEMENTATION

Nếu không được truyền trong mã, sp_size có thể được đặt thông qua đối số CLI --parallelism_config_sp_size. Tương tự cho các đối số khác. Bạn cũng có thể sử dụng cấu hình theo kiểu tệp accelerate config, ví dụ, cho 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

Như đã đề cập ở trên, song song hóa chuỗi Ulysses thường được kết hợp với song song hóa dữ liệu - cùng các rank được sử dụng để cung cấp các luồng dữ liệu riêng biệt và đồng thời thực hiện Song song hóa Chuỗi Ulysses. Tuy nhiên, bạn cũng có thể tạo các bản sao như sau:

# 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(...),
)

Ở đây chúng ta sử dụng 4 GPU, với 2 bản sao song song chuỗi (sequence parallelism). Deepspeed-ZeRO là thành phần điều phối song song dữ liệu (data parallelism) ở đây.

Lưu ý rằng có rất nhiều cơ chế phức tạp được ẩn bên trong UlyssesSPDataLoaderAdapter. Thành phần này hoạt động ở phía sau, bao bọc đối tượng DataLoader gốc của bạn, nhưng bạn nên nắm rõ điều này trong trường hợp gặp sự cố. Nó cũng tự động chèn shift_labels chính xác vào từ điển batch, trước khi batch được chia nhỏ trên các rank tham gia.

Bây giờ, phần duy nhất còn lại để bắt đầu sử dụng ALST/UlyssesSP là tổng hợp loss giữa các rank bằng cách sử dụng all_gather có thể tính đạo hàm để lấy đúng các grads. Đoạn mã sau thực hiện điều này, đồng thời loại bỏ các token bị mask bằng -100, để có được giá trị trung bình chính xác:

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

Lưu ý rằng các model tự động xử lý shift_labels khi tham số này có mặt trong batch. Forward pass của model nhận shift_labels thông qua **kwargs và truyền nó đến hàm loss, từ đó tính toán loss chính xác cho sequence parallelism. Nếu bạn sử dụng Liger Kernel, nó cũng xử lý shift_labels một cách liền mạch và tính toán loss với hiệu quả bộ nhớ rất cao. Liger được khuyến nghị mạnh mẽ cho các độ dài sequence dài, vì nó giải phóng bộ nhớ GPU bằng cách sử dụng các thao tác đã được hợp nhất (ví dụ: tính toán logit-loss hợp nhất không bao giờ tạo ra tensor logits đầy đủ trong bộ nhớ).

Nếu bạn muốn xem HF Accelerate đã làm gì phía sau, hãy đọc tutorial tích hợp đầy đủ này.

Để xem một ví dụ về vòng lặp huấn luyện Accelerate với ALST/UlyssesSP được bật, hãy xem examples/alst_ulysses_sequence_parallelism.

[!Cảnh báo]

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, một số phần nhỏ của API công khai có thể thay đổi trong tương lai.

Vì đây là backend Deepspeed nên cấu hình Deepspeed thông thường được áp dụng, do đó bạn có thể kết hợp sequence parallelism với việc offload trạng thái optimizer và/hoặc trọng số để giải phóng thêm bộ nhớ GPU và cho phép độ dài sequence dài hơn nữa. Công nghệ này đã được kiểm chứng hoạt động với DeepSpeed ZeRO stage 2 và 3.