Context Parallel 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 context parallelism trong 🤗accelerate, dành cho những độc giả tò mò hơn, chúng tôi cũng sẽ đề cập đến một số chi tiết kỹ thuật trong các phần sau.

Xem thêm Hướng dẫn về Song song Chuỗi rất liên quan.

Tại sao cần song song ngữ cảnh?

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ể.

Context 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 song song trên nhiều GPU. Nhờ đó, chúng ta có thể huấn luyện các mô hình với sequence dài, mở rộng quy mô lên đến độ dài sequence 1M+.

Cách sử dụng song song ngữ cảnh?

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

Tương tự như mọi tính năng khác trong 🤗accelerate, bạn cũng có thể bật context parallelism 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-cp-size 8 --parallelism-config-cp-comm-strategy [allgather|alltoall] ...

Bạn cũng có thể đặt cp_size và cp_comm_strategy 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.

Context parallelism tương thích với các chiến lược song song khác, chẳng hạn như data parallelism, tensor parallelism và FSDP2. Bạn có thể kết hợp chúng một cách đơn giản bằng cách đặt kích thước song song của bạn thành các giá trị mong muốn, ví dụ --parallelism-config-dp-size 8 --parallelism-config-tp-size 2 --parallelism-config-cp-size 8. Hoặc bạn có thể sử dụng lớp ParallelismConfig để đặt chúng theo chương trình.

Context parallelism được ghép nối chặt chẽ với FSDP2, mà bạn có thể tìm hiểu thêm trong giới thiệu về FSDP2. Điều này có nghĩa là context parallelism chỉ hoạt động nếu bạn sử dụng FullyShardedDataParallelPlugin hoặc --use-fsdp với version được đặt thành 2 trong chương trình của bạn. Nếu không sử dụng FSDP2, một lỗi sẽ được ném ra.

Context parallelism chỉ hoạt động với SDPA và chỉ khi không có mask hoặc causal mask. Chúng tôi không thể phát hiện điều này một cách chính xác cho bạn, vì vậy trách nhiệm của bạn là đảm bảo rằng bạn đang sử dụng SDPA mà không có mask hoặc causal mask. Nếu bạn sử dụng bất kỳ triển khai attention nào khác, một lỗi sẽ được ném ra.

Sau khi bật context parallelism bằng các phương pháp được đề cập ở trên, bạn có thể áp dụng nó vào vòng lặp huấn luyện của mình. Chúng tôi cung cấp một lớp bọc mỏng quanh torch.distributed.tensor.experimental.context_parallel mà bạn có thể sử dụng trong vòng lặp huấn luyện, giúp trừu tượng hóa một phần độ phức tạp khi sử dụng nó (chi tiết hơn sẽ được trình bày sau). Để giảm thiểu các thay đổi bạn phải thực hiện trong vòng lặp huấn luyện, chúng tôi cung cấp một context manager, vốn là noop nếu context parallelism không được bật, và sẽ áp dụng context parallelism nếu nó được bật. Nhờ đó, bạn có thể sử dụng nó trong vòng lặp huấn luyện mà không cần thay đổi bất kỳ mã nào dựa trên cấu hình song song của mình. Bạn có thể sử dụng nó như sau:

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 này phải được tạo lại ở mỗi bước huấn luyện, như minh họa trong ví dụ ở trên. Việc này là rất quan trọng.

Điều này có thể mở rộng kích thước ngữ cảnh của bạn lên độ dài chuỗi 1M+. Bên dưới, chúng tôi trình bày tốc độ và mức sử dụng bộ nhớ của context parallelism cho kích thước ngữ cảnh lên đến 256k. Chúng ta có thể thấy rằng khi nhân đôi kích thước ngữ cảnh và số lượng GPU, chúng ta có thể đạt được mức sử dụng bộ nhớ ổn định, từ đó có khả năng cho phép mở rộng độ dài ngữ cảnh không giới hạn.

Hình 1: Mức sử dụng bộ nhớ và tốc độ của context parallelism với kích thước context lên đến 256k.

Các ví dụ này được tạo bằng một script mà bạn có thể tìm thấy trong thư mục examples. Để chạy ví dụ trên 8 GPU H100 (độ dài chuỗi 128k), bạn có thể sử dụng lệnh sau:

Đã sao chép

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

Giao diện của Accelerate

Trình quản lý ngữ cảnh nhận một vài đối số, được dùng để cấu hình song song ngữ cảnh.

• buffers: Đây là danh sách các tensor sẽ được chia nhỏ (shard) theo chiều sequence. Các tensor này thường là input ids, labels và attention mask.
• buffer_seq_dims: Đây là một danh sách các số nguyên, chỉ định chiều sequence của các buffer, theo thứ tự của danh sách buffers. Nếu bạn truyền buffers=[input_ids, shift_labels] với cả hai đều có shape [batch_size, sequence_length], bạn sẽ truyền buffer_seq_dims=[1, 1]. vì chiều sequence là chiều thứ hai của các tensor. Điều này là bắt buộc để tính toán chính xác các đầu ra của model.
• no_restore_buffers: Việc triển khai context parallelism sẽ sửa đổi các buffer tại chỗ, chuyển đổi chúng thành torch.distributed.tensor.Dtensor. Sau khi context manager thoát, một kernel giao tiếp sẽ cần được khởi chạy để khôi phục các buffer về trạng thái ban đầu (thường là all-gather). Quá trình này mất một ít thời gian, vì vậy được khuyến nghị truyền các tensor giống như trong đối số buffers, để tránh giao tiếp không cần thiết, trừ khi bạn chắc chắn cần sử dụng các buffer sau khi context manager thoát.

Context parallelism không tương thích với labels là bản sao của input_ids, vốn các model từ 🤗 transformers có thể tự dịch chuyển để kích hoạt causal language modeling. Hãy tưởng tượng trường hợp này: labels = [l1, l2, l3, l4, … li] nếu chúng ta áp dụng context parallelism, mỗi rank sẽ nhận được một phần của labels, chẳng hạn như: labels_rank_0 = [l1, l2], labels_rank_1 = [l3, l4], … sau khi mã hóa modelling của transformers dịch chuyển labels, kết quả sẽ là: labels_rank_0 = [l2, PAD], labels_rank_1 = [l3, PAD], … trong đó PAD là token padding. Điều này sẽ dẫn đến việc tính toán loss không chính xác, vì labels không còn được căn chỉnh với inputs nữa. Vì lý do này, bạn cần tự động dịch chuyển labels trước khi truyền chúng vào model

Các tùy chọn có thể cấu hình

Accelerate chỉ cung cấp một tùy chọn duy nhất để cấu hình song song ngữ cảnh (ngoại trừ cp_size)

• cp_comm_strategy: Phương pháp xoay vòng để sử dụng cho các shard. Chúng tôi khuyến nghị mạnh mẽ giữ nguyên giá trị "allgather", vì trong hầu hết các trường hợp, nó rất có khả năng cho hiệu suất vượt trội hơn so với "alltoall".

Context parallel size khá tự giải thích, đó là số lượng rank mà các đầu vào sẽ được chia nhỏ (sharded) trên đó. Context parallel shard rotation xác định cách các shard của đầu vào được xoay vòng trên các rank. Chúng ta sẽ xem xét chi tiết 2 tùy chọn này trong phần tiếp theo.

Bạn có thể xem một ví dụ end-to-end trong tệp ND parallel example, nơi bạn có thể huấn luyện một mô hình 8B với độ dài context lên đến 128k trên một node 8xH100 duy nhất. Sử dụng huấn luyện đa node, bạn có thể mở rộng quy mô này lên độ dài sequence 1M+ trên nhiều GPU. Bạn cũng có thể kết hợp liền mạch với các chiến lược song song khác để phù hợp với nhu cầu của mình.

Chi tiết kỹ thuật

Phần này khá mang tính kỹ thuật, vì vậy nếu bạn không cần tìm hiểu chi tiết bên trong của context parallelism, bạn có thể bỏ qua và bắt đầu xây dựng 🚀

Chúng ta sẽ sử dụng từ shard rất nhiều trong các phần tiếp theo, vì vậy hãy định nghĩa nó trước. Nếu chúng ta gọi một tensor là sharded trên chiều Dth, trên N rank, điều đó có nghĩa là tensor này được chia thành N phần, trong đó mỗi phần của tensor có shape [..., D//N, ...].

Vậy nó hoạt động như thế nào?

Context parallelism hoạt động bằng cách chia nhỏ các ma trận Q, K and V theo chiều sequence. Mỗi rank có shard Q được phân công, gọi là Q_i. Ma trận này chỉ tồn tại trên rank này trong suốt quá trình tính toán. Tương tự, mỗi rank có shard riêng của K và V, gọi là K_i và V_i. Sau đó, mỗi rank tính attention với shard riêng của Q_i, K_i và V_i, gọi là attn_i. Trong quá trình tính toán này, một communication kernel được khởi chạy để thu thập các Ks và Vs từ tất cả các rank khác. Primitive giao tiếp được sử dụng phụ thuộc vào tùy chọn context_parallel_shard_rotation. Nhờ đó, mỗi rank có thể tính local attention, trước tiên với Q_i, K_i và V_i, sau đó với K_j và V_j từ tất cả các rank khác. Vì mỗi rank giữ các ma trận Q, K and V được chia nhỏ theo chiều sequence, các ma trận kết quả sẽ nhỏ hơn và có thể vừa trên một GPU duy nhất.

Chúng ta có thể hình thức hóa điều này bằng đoạn mã giả sau:

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 so với all-gather

all-gather

Vậy sự khác biệt giữa all-to-all và all-gather là gì? Với all-gather, quá trình truyền thông rất đơn giản. Sau khi (thực ra là trước, vì thường mất nhiều thời gian hơn) chúng ta tính toán local attention attn_i, chúng ta khởi chạy all-gather để thu thập tất cả các Ks và Vs từ tất cả các rank khác. Khi quá trình truyền thông này hoàn tất, mỗi rank sẽ có tất cả các Ks và Vs từ các rank khác và có thể tính toán attention với chúng một cách tuần tự. Trong kịch bản lý tưởng, all-gather sẽ hoàn thành đúng vào thời điểm tính toán attn_i kết thúc. Tuy nhiên, điều này không bao giờ xảy ra trong thực tế, nên mức độ chồng chéo lý tưởng thực sự đạt được khi toàn bộ attn_i được chồng chéo với một phần của quá trình truyền thông, sau đó để bắt đầu tính toán với K_j và V_j, chúng ta chờ cho đến khi all-gather hoàn tất.

all-to-all

All-to-all, hay đôi khi được gọi là ring-rotation, sử dụng một mẫu giao tiếp dạng vòng. Sau khi hoàn tất tính toán attn_i, một all-to-all được khởi chạy để gửi K_i và V_i đến các rank lân cận. Sau đó, chúng ta lặp lại quá trình này context_parallel_size-1 lần, để mỗi rank nhìn thấy tất cả các shard của K và V từ tất cả các rank khác một lần. Trong kịch bản lý tưởng, chúng ta prefetch các shard K_i+1 và V_i+1 từ rank lân cận và giao tiếp này được che phủ chính xác với việc tính toán attn_i hiện tại của chúng ta. Một lần nữa, trên thực tế, sự che phủ hoàn hảo này không bao giờ xảy ra. Do bản chất của cách tiếp cận này, nếu chúng ta không đạt được sự che phủ hoàn hảo, mức phạt sẽ lớn hơn rất nhiều so với all-gather.

Cách chọn phương pháp xoay phù hợp?

Về mặt lý thuyết, all-to-all nên là lựa chọn tốt hơn. Tuy nhiên, trên thực tế, điều này hiếm khi xảy ra. Do đó, mặc định chúng tôi sử dụng all-gather, vì phương pháp này có khả năng đạt hiệu suất tốt hơn cao hơn. Các benchmark chuyên sâu từ nhóm torchtitan cũng cho thấy all-to-all hiếm khi vượt trội hơn all-gather. Mặc dù vậy, chúng tôi vẫn cung cấp cả hai tùy chọn, vì bạn có thể thấy một trong số chúng phù hợp hơn với trường hợp sử dụng của mình.

Bạn có thể trực tiếp nhìn thấy vấn đề này trong đầu ra của trình phân tích hiệu năng ở hình ảnh bên dưới:

Hình 1: Phần màu đỏ cho thấy thời gian chờ nhàn rỗi trong khi chúng ta chờ kernel all-to-all hoàn tất. Thanh màu xanh đầu tiên được tô sáng cho thấy mất ~250us để hoàn tất, và điều này được lặp lại N-1 lần cho mỗi lần gọi attention, trong đó N là kích thước context parallel.

Tại sao chỉ có FSDP2?

Chúng tôi chỉ hỗ trợ context parallelism với FSDP2, vì chúng tôi tạo một joint mesh của context_parallel_size và dp_shard_size để khai thác tối đa tiềm năng của nó. Cách hoạt động là: chúng tôi shard mô hình trên joint mesh có kích thước cp_size*dp_shard_size, điều này tối đa hóa việc tiết kiệm bộ nhớ. Đây là một dạng “free lunch”, vì giao tiếp của FSDP được overlap hoàn toàn với việc tính toán attention, như được minh họa trong các hình ảnh bên dưới.

Hình 2: Trong các hình chữ nhật màu xanh (Stream 23), bạn có thể thấy rằng việc pre-fetch shard của `FSDP` được chồng hoàn toàn lên phần tính toán attention (Stream 7), trong khi ở các hình chữ nhật màu đỏ (Stream 24), bạn có thể thấy rằng kết quả của kernel all-gather tạo ra một khoảng thời gian rảnh, trong đó stream tính toán của chúng ta (7) đang ở trạng thái rảnh.

Trong hình trên, bạn cũng có thể nhận thấy sự khác biệt giữa all-to-all và all-gather. Trong khi ở all-to-all (Hình 1), chúng ta khởi chạy một kernel truyền thông N-1 lần cho mỗi lần gọi attention, thì ở all-gather (Hình 2), chúng ta chỉ khởi chạy một kernel truyền thông duy nhất. Điều này dẫn đến một bong bóng lớn hơn, nhưng nó chỉ xảy ra một lần cho mỗi lần gọi attention, trong khi ở all-to-all, nó xảy ra N-1 lần.

Phân phối dữ liệu trong lưới kết hợp

Chúng tôi đảm bảo gửi cùng một lô dữ liệu đến toàn bộ nhóm con cp, để kết quả chính xác. (Nghĩa là mỗi rank trong nhóm con cp nhận cùng một lô dữ liệu.) Tuy nhiên, chúng tôi cũng gửi các lô dữ liệu khác nhau đến từng rank của nhóm dp_shard. Hãy hình dung như sau:

# 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.