Song song chuỗi Ulysses

Ulysses sequence parallelism (SP) huấn luyện trên các chuỗi rất dài bằng cách chia chúng trên nhiều GPU. Để tính toán attention chính xác, một tập thể all-to-all sẽ hoán đổi chiều sharding từ sequence sang attention heads. Mỗi GPU sau đó có toàn bộ sequence và tính toán attention cục bộ trên một tập con các heads. Một all-to-all thứ hai đưa về bố cục sharding theo sequence để phần còn lại của forward pass tiếp tục cục bộ trên mỗi chunk.

                        GPU 0                       GPU 1
                   ┌───────────────┐          ┌───────────────┐
  forward          │ tokens 0..N/2 │          │ tokens N/2..N │  ← each GPU holds half the sequence
  (seq-sharded)    │  all H heads  │          │  all H heads  │
                   └───────┬───────┘          └───────┬───────┘
                           └───────── all-to-all ──────┘
                   ┌───────────────┐          ┌───────────────┐
  attention        │ all N tokens  │          │ all N tokens  │  ← now each GPU has the full sequence
  (head-sharded)   │ heads 0..H/2  │          │ heads H/2..H  │  ← but only half the heads
                   └───────┬───────┘          └───────┬───────┘
                           └───────── all-to-all ──────┘
                   ┌───────────────┐          ┌───────────────┐
  forward          │ tokens 0..N/2 │          │ tokens N/2..N │  ← back to seq-sharded
  (seq-sharded)    │  all H heads  │          │  all H heads  │
                   └───────────────┘          └───────────────┘

Hướng dẫn này bao gồm thành phần song song hóa chuỗi Ulysses của ALST (Arctic Long Sequence Training). Hệ thống ALST đầy đủ cũng bao gồm TiledMLP và offloading checkpoint kích hoạt, những tính năng không có sẵn trong Transformers. Xem DeepSpeed ALST tutorial để biết về hệ thống hoàn chỉnh.

Cấu hình

Sequence parallelism yêu cầu Accelerate v1.12.0 và ít nhất 2 GPU. Cấu hình sequence parallelism trong ParallelismConfig của Accelerate và truyền nó vào parallelism_config hoặc một tệp cấu hình Accelerate.

parallelism_config

Tệp cấu hình Accelerate

Các trường sau đây quan trọng để cấu hình song song chuỗi.

Trainer tự động xử lý việc chia nhỏ DataLoader, tạo position_ids, dịch chuyển nhãn và tổng hợp loss giữa các rank SP. Nếu bạn đang viết vòng lặp huấn luyện tùy chỉnh, hãy xem hướng dẫn Sequence Parallelism của Accelerate.

sp_backend phải được đặt thành "deepspeed" để sử dụng Ulysses sequence parallelism.
sp_size là số lượng GPU xử lý một chuỗi duy nhất song song. Mỗi rank SP nhận một luồng dữ liệu riêng biệt từ DataLoader, khác với tensor parallelism, nơi tất cả các rank nhận cùng một dữ liệu. dp_world_size = world_size / sp_size hiệu dụng, do đó với 4 GPU và sp_size=4, dp_world_size=1 cho các phép tính batch size. Các chuỗi cũng phải được đệm (padding) đến bội số của sp_size. Hãy đặt pad_to_multiple_of trong data collator của bạn cho phù hợp.

Số lượng đầu chú ý phải chia hết cho sp_size. Một mô hình có 32 đầu hỗ trợ sp_size là 1, 2, 4, 8, 16 hoặc 32.
```
from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
    pad_to_multiple_of=sp_size,
)
```
sp_seq_length_is_variable kiểm soát việc xử lý độ dài chuỗi thay đổi. Đặt thành True (khuyến nghị) khi độ dài thay đổi giữa các batch. Đặt thành False khi tất cả các chuỗi được pad về một độ dài cố định được chỉ định bởi sp_seq_length.
sp_attn_implementation thiết lập backend cho attention. Các giá trị được hỗ trợ là "sdpa", "flash_attention_2", hoặc "flash_attention_3". FlashAttention được khuyến nghị, đặc biệt khi đóng gói nhiều mẫu trong một batch. SDPA có thể xử lý attention sai lệch qua các ranh giới mẫu khi các mẫu được đóng gói. Eager attention không được hỗ trợ vì attention_mask 4D của nó bị loại bỏ do lý do bộ nhớ và khả năng mở rộng.

Kết hợp với song song hóa dữ liệu

Sequence parallelism và data parallelism sử dụng cùng các GPU, và SP không yêu cầu phần cứng bổ sung. Để chạy cả hai, hãy đặt dp_replicate_size hoặc dp_shard_size sao cho dp_replicate_size × dp_shard_size × sp_size bằng tổng số GPU của bạn.

Ví dụ, với 8 GPU và sp_size=4, hãy đặt dp_replicate_size=2 (2 × 1 × 4 = 8).

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_replicate_size=2,
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2",
    ),
)

Các bước tiếp theo

Song song ngữ cảnh với FSDP2 (cp_size) là tách biệt với Ulysses (sp_size). Xem Đánh giá với song song ngữ cảnh.
Hướng dẫn Sequence Parallelism của Accelerate đi sâu hơn vào việc triển khai Ulysses và cho thấy cách viết một vòng lặp huấn luyện tùy chỉnh.
Hướng dẫn sử dụng DeepSpeed ALST bao quát toàn bộ hệ thống ALST, bao gồm TiledMLP và việc offload checkpoint kích hoạt.
Hướng dẫn về các phương pháp song song hóa cho thấy cách kết hợp song song hóa chuỗi với các chiến lược khác như ZeRO.
Bài đăng trên blog Ulysses Sequence Parallelism: Training with Million-Token Contexts giải thích cách Ulysses hoạt động và cách nó được tích hợp trong Accelerate, Trainer và SFTTrainer.

Update on GitHub