Song song chuỗi siêu dài- Nguyên tắc Kỹ thuật và Triển khai Ulysses + Ring-Attention

November 10, 2025
15 min read

Siêu phân tán chuỗi dài: Các nguyên tắc kỹ thuật và triển khai của Ulysses + Ring-Attention

Đào tạo chuỗi siêu dài luôn là một hướng đi quan trọng trong việc đào tạo các mô hình lớn. Trong các quy trình suy luận thực tế, đặc biệt là trong quy trình hoạt động của Agent, khả năng tổng quát hóa của mô hình đối với các chuỗi dài và các tình huống phức tạp đại diện cho độ tin cậy của mô hình trong các ứng dụng thực tế. Các tình huống chuỗi dài cũng đặt ra yêu cầu cao hơn đối với việc đào tạo các mô hình lớn. Do đặc điểm độ phức tạp O(N^2) của tính toán Attention, khi chuỗi đầu vào thực tế tăng lên, việc sử dụng bộ nhớ có sự bùng nổ theo cấp số nhân. Điều này đặt ra những thách thức to lớn cho khả năng sử dụng các loại card có bộ nhớ hạn chế trong các tình huống đào tạo chuỗi dài.

Công nghệ Sequence Parallel (SP) có thể được sử dụng để giảm sự phụ thuộc vào bộ nhớ lớn cho việc đào tạo chuỗi dài trong điều kiện đa card hoặc đa máy. Nói một cách đơn giản, song song hóa chuỗi có thể được định nghĩa bởi khái niệm sau:

Song song hóa chuỗi là quá trình chia một chuỗi đầu vào thành nhiều chuỗi con để tính toán song song trên các card khác nhau trong quá trình đào tạo, do đó giảm yêu cầu bộ nhớ cho việc đào tạo.

Các phương pháp song song hóa chuỗi phổ biến bao gồm:

Ulysses
Ring-Attention
Megatron-SP/CP

Trong số này, cả Ulysses và Ring-Attention đều là các giải pháp song song hóa chuỗi dựa trên hệ sinh thái Transformers. Chúng tôi chủ yếu giới thiệu các nguyên tắc kỹ thuật và triển khai của hai giải pháp này. Megatron-CP[1] và Ring-Attention có thể được coi là các công nghệ tương tự, trong khi Megatron-SP thực hiện phân chia trên các giá trị kích hoạt và thường được sử dụng kết hợp với Megatron-TP. Chúng tôi sẽ không đi sâu vào những điều này ở đây.

Trước hết, hãy xem xét hiệu quả giảm bộ nhớ có thể đạt được trong việc đào tạo chuỗi dài trên mô hình Qwen2.5-3B sau khi tích hợp cả hai công nghệ đào tạo Ulysses và Ring-Attention:

Kích thước SP	Chiến lược SP	Bộ nhớ GPU	Thời gian đào tạo
8	w/ ulysses=2 ring=4	17.92GiB	1:07:20
4	w/ ulysses=2 ring=2	27.78GiB	37:48
2	w/ ulysses=2	48.5GiB	24:16
1	w/o SP	75.35GiB	19:41

Lưu ý rằng do khối lượng giao tiếp tăng lên và tải GPU khác nhau gây ra bởi việc chia chuỗi, thời gian đào tạo sẽ tương ứng kéo dài.

Khi chia thành 2 chuỗi con, Ulysses đã được sử dụng; khi chia thành 4/8 chuỗi con, Ulysses (chia 2) + Ring-Attention (chia 2 hoặc 4) đã được sử dụng. Dưới đây chúng tôi đi sâu vào các nguyên tắc và kế hoạch triển khai của hai công nghệ này.

Ulysses

Ulysses là một thuật toán song song hóa chuỗi được phát triển bởi nhóm DeepSpeed[2]. Ý tưởng của Ulysses có thể được tóm tắt trong một câu:

Sau khi chuỗi được chia thành các chuỗi con, việc trao đổi giá trị kích hoạt được thực hiện trước tính toán Attention trong mỗi lớp, cho phép mỗi card kết hợp thành một chuỗi hoàn chỉnh. Sau đó, các Attention Heads được phân phối trên các card khác nhau để đạt được việc giảm bộ nhớ. Sau khi tính toán hoàn tất, chúng sẽ được trao đổi trở lại.

Thông qua cách tiếp cận này, mặc dù tính toán QKV vẫn liên quan đến các giá trị kích hoạt O(N^2), bộ nhớ vẫn được giảm vì mỗi card có ít Attention Heads hơn. Hình sau đây cho thấy một ví dụ đơn giản, giả sử chia thành hai chuỗi, mỗi chuỗi có hai attention heads:

Đây là hình ảnh từ bài báo:

Trong đó N biểu thị độ dài chuỗi và d biểu thị hidden_size, cũng có thể được hiểu cụ thể là số lượng Attention Heads * hidden_size thực tế.

Trạng thái sau khi chia Ulysses:

Nguyên tắc kỹ thuật của Ulysses rất rõ ràng, với điểm mấu chốt là giao tiếp all-to-all từ N/P đến d/P. Vì QKV hoàn chỉnh trong quá trình tính toán Attention, và các Attention-Heads khác nhau được phân phối trên các card khác nhau, nó có tính phổ quát trong các tình huống như GQA và MHA, và hoàn toàn tương thích với nhiều công nghệ khác nhau như flash-attn, SDPA và padding_free. Tất nhiên, do sự tồn tại của giao tiếp giữa các card, quá trình backward yêu cầu xử lý thêm.

Tuy nhiên, những hạn chế của Ulysses cũng khá rõ ràng, đó là bị giới hạn bởi số lượng Attention Heads. Đặc biệt trong GQA, nơi số lượng KV heads ít hơn nhiều so với số lượng Q heads, Ulysses có thể không thể chia trên nhiều card hơn.

Ring-Attention

Trước khi thảo luận về Ring-Attention, chúng ta cần nói ngắn gọn về Flash-Attention. Nguyên tắc của Flash-Attention cũng có thể được giải thích trong một câu:

QKV và softmax có thể thực hiện tính toán và cập nhật song song theo khối, tối đa hóa việc sử dụng SRAM đồng thời giảm việc sử dụng bộ nhớ.

Mã giả cho quá trình chuyển tiếp của flash-attention:

Lưu ý mã giả ở dòng 10 ~ 12 của Thuật toán 1 ở trên, phần này thực hiện cập nhật hợp nhất trên * LSE* (log-sum-exp) theo khối:

$$ lse_i^{new} = m_i^{new} + \log\left(e^{m_i - m_i^{new}}\ell_i + e^{\tilde m_{ij}-m_i^{new}} \tilde \ell_{ij}\right) $$

Trong đó:

$$ m_i^{new} = \max(m_i, \tilde m_{ij}) $$

Và Attention-Out được hợp nhất và cập nhật, tức là sau khi tính toán các khối mới, kết quả của các khối mới được hợp nhất với kết quả của các khối cũ để thu được kết quả cuối cùng.

Vì vậy, nếu mỗi card mang một phần của độ dài chuỗi và kết quả tính toán được truyền qua các card, flash-attention có thể hoạt động giữa các card không? Đây là ý tưởng cơ bản của Ring-Attention.

Ring-Attention: Tận dụng nguyên tắc rằng tính toán Attention có thể được thực hiện theo khối, các khối chuỗi được chia trên nhiều card để tính toán riêng biệt, và sau đó kết quả tính toán được hợp nhất để thu được kết quả cuối cùng.

Giả sử có N khối, hãy xem xét cùng Q$i$ và K${0..n-1}$~ V$_{0..n-1}$~ có thể giao tiếp và chảy giữa các khối khác nhau. Trước tiên hãy xem xét phần Softmax:

$$ p_{ij} = \frac{e^{x_{ij}}}{Z_i} $$

Trong đó:

$$ Z_i = \sum_{j=1}^N e^{x_{ij}} $$

Thông thường, để tránh bất ổn định số học, số mũ không được tính toán trực tiếp, mà sử dụng các phương pháp tính toán ổn định số học:

$$ p_{ij} = \exp(x_{ij} - \text{lse}_i) $$

$$ \text{lse}_i = \log \sum_{j=1}^N e^{x_{ij}} $$

Có thể thấy rằng sau khi mở rộng bằng công thức số mũ, phương pháp này tương đương với công thức ban đầu ở trên.

Tiếp theo, chúng ta cần suy ra các bản cập nhật cho LSE và Attention-Out một cách đệ quy. Hãy xem xét LSE trước. Xem xét chúng ta đã có kết quả tích lũy trước đó và kết quả khối hiện tại, LSE mới sẽ là:

$$ Z_i^{new} = \sum_{j\in\text{prev}} e^{x_{ij}} + \sum_{j\in\text{block}} e^{x_{ij}} $$

Sau đó:

$$ A_i^{new} = A_i + \tilde A_{ij},\quad Z_i^{new} = Z_i + \tilde Z_{ij},\quad {out}_i^{new} = \frac{A_i^{new}}{Z_i^{new}} $$

Sử dụng LSE để biểu thị công thức trên:

$$ {out}_i^{new} = \frac{A_i+ \tilde A_{ij}}{e^{\text{lse}_i^{new}}} $$

Lưu ý rằng sau khi chia tử số trong công thức trên, nó có thể được chuyển đổi thành tổng của hai biểu thức độc lập, và hai biểu thức này tương ứng ở dạng sigmoid, do đó chúng ta thu được:

$$ {out}_i^{new} = \text{sigmoid}(\text{lse}_i - \tilde{\text{lse}}_{ij}) \cdot \text{out}_i + \text{sigmoid}(\tilde{\text{lse}}_{ij} - \text{lse}_i) \cdot \tilde{\text{out}}_{ij} $$

Hai công thức cập nhật này tương đương với các công thức đệ quy được đưa ra trong bài báo flash-attention, cả hai đều tuân theo phương pháp cập nhật khối và online-softmax.

Quá trình suy luận trên cung cấp công thức tính toán chuyển tiếp cho các cập nhật lặp lại, với mã nằm trong update_out_and_lse: https://github.com/modelscope/ms-swift/blob/main/swift/trainers/sequence_parallel/zigzag_ring_attn.py#L69

Vì có một quá trình chuyển tiếp, chắc chắn phải có một quá trình ngược. Trong quá trình lan truyền ngược, chúng ta cần dần dần khôi phục từ bước i-1 cuối cùng về bước 0. Do hạn chế về không gian, việc suy luận công thức ngược không được mở rộng ở đây, với mã nằm trong lse_grad: https://github.com/modelscope/ms-swift/blob/main/swift/trainers/sequence_parallel/zigzag_ring_attn.py#L263 https://github.com/modelscope/ms-swift/blob/main/swift/trainers/sequence_parallel/zigzag_ring_attn.py#L458

Tốt, chúng ta đã chuẩn bị lý thuyết ở trên và có thể bắt đầu triển khai mã.

Lưu ý rằng Ring-Attention có nhiều biến thể triển khai, chẳng hạn như strip-ring-attention[3]. Trong số các triển khai này, triển khai zigzag là xuất sắc nhất về cân bằng tải. Để hiểu vấn đề với Ring-Attention ban đầu, hãy xem hình sau:

Do GPU 0 xử lý phần đầu của câu, khi KV từ các card khác chảy đến GPU 0, GPU không thể tham gia tính toán phần sau của câu do causal=True, trong khi GPU 3 có thể tính toán toàn bộ chuỗi từ 0 đến 2. Do đó, tải tính toán trên mỗi card không nhất quán. Dưới tiền đề này, Megatron-CP và một số triển khai xuất sắc áp dụng phân chia zigzag:

Giả sử chúng ta cần phân chia trên 4 card, sau đó trong điều kiện chuỗi có thể được chia đều thành 8 phần, chúng ta kết hợp 0/7 với nhau, 1/6 với nhau, 2/5 và 3/4 tương ứng với nhau. Điều này đảm bảo tính toán cân bằng. Hơn nữa, tính toán này còn có một đặc điểm khác:

Khi tính toán QKV cục bộ (số chuỗi 0), causal=True được tính toán trực tiếp
Khi số chuỗi chảy nhỏ hơn hoặc bằng thứ hạng hiện tại, chỉ cần tính toán nửa đầu của KV
Khi số chuỗi chảy lớn hơn thứ hạng hiện tại, chỉ cần tính toán nửa sau của Q

Điều này càng làm giảm tải tính toán. Để triển khai mã, vui lòng xem: https://github.com/modelscope/ms-swift/blob/main/swift/trainers/sequence_parallel/zigzag_ring_attn.py#L348

Ulysses và Ring-Attention

Không khó để thấy rằng hai sơ đồ song song hóa chuỗi này đều có những đặc điểm riêng.

Ulysses có chi phí giao tiếp tương đối thấp, nhưng bị giới hạn bởi số lượng Attention Heads, và giao tiếp all-to-all nhạy cảm với độ trễ và có các yêu cầu nhất định về cấu trúc liên kết mạng.
Giao tiếp vòng P2P của Ring-Attention có yêu cầu thấp hơn, nhưng có khối lượng giao tiếp cao hơn và không bị giới hạn bởi số lượng Attention Heads.

Từ các nguyên tắc trên, chúng ta có thể thấy rằng Ulysses và Ring-Attention thực sự có thể được sử dụng kết hợp. Chúng ta có thể sử dụng Ulysses với khối lượng giao tiếp thấp hơn để phân chia trước, và nếu số lượng Attention Heads không đủ (GQA), hoặc số lượng chuỗi được phân chia quá lớn, thì bổ sung bằng Ring-Attention.

SWIFT triển khai công nghệ tính toán kết hợp như vậy, có thể áp dụng cho nhiều tình huống khác nhau bao gồm văn bản thuần túy, đa phương thức, SFT, DPO, GRPO, v.v. Trong quá trình triển khai mã cơ bản, chúng tôi đã áp dụng một số tác phẩm mã nguồn mở xuất sắc của cộng đồng[4][5] và viết lại một phần mã.

https://github.com/modelscope/ms-swift

Cách sử dụng cũng rất đơn giản, chỉ cần thêm một tham số bổ sung vào dòng lệnh:

bash –sequence_parallel_size N

Framework sẽ tự động tính toán phương pháp phân chia, và thậm chí có thể hỗ trợ các trường hợp số lượng GPU không đều (3, 5, 7, v.v.).

Phương pháp phân chia

Cách tiếp cận tự nhiên nhất là sử dụng Ulysses để thu thập cục bộ trước, sau đó sử dụng Ring-Attention toàn cục để tính toán LSE và Attention-Out toàn cục. Giả sử phân chia thành 4 chuỗi con (Ulysses world_size=2, Ring-Attention world_size=2), với model head=4:

Sau khi giao tiếp all-to-all của Ulysses, GPU0 và GPU1 như cùng một nhóm Ulysses đều giữ chuỗi 0/3, nhưng với các head khác nhau (nửa đầu và nửa sau). GPU2 và GPU3 cũng tương tự. Trong quá trình tính toán Ring-Attention, GPU0 và GPU2 tạo thành một nhóm Ring-Attention để giao tiếp vòng, và GPU1 và GPU3 cũng tương tự.

Trước khi phân chia, chuỗi cần được đệm để nó có thể được chia cho world_size*2 (nhân 2 vì zigzag yêu cầu kết hợp lại các khối con).

Thích ứng đa phương thức

Thích ứng phân chia chuỗi cho các mô hình đa phương thức khá khó khăn, chủ yếu là do:

Độ dài chuỗi của các mô hình đa phương thức không thể xác định trước khi thực hiện chuyển tiếp thực tế. Một số mô hình chỉ sử dụng một thẻ <image> token để đại diện cho phần đa phương thức, và sau khi ViT mã hóa hình ảnh, thẻ này được thay thế bằng một chuỗi rất dài.
Một số chuỗi đầu vào của mô hình chứa các thẻ đóng, chẳng hạn như <image></image>, không thể được phân chia trước khi chúng được thay thế bằng mã hóa hình ảnh thực tế, nếu không lỗi sẽ được ném ra trực tiếp.

Nói chung, LLM đa phương thức chứa các lớp bên trong và bên ngoài của mô hình. Lớp ngoài bao gồm xử lý ViT và logic tính toán lm_head, trong khi lớp bên trong tính toán decode_layers, mà chúng ta gọi là backbone.

Để thích ứng phân chia đa phương thức, SWIFT áp dụng một mẹo kỹ thuật trong quá trình triển khai: phân chia không xảy ra trong quá trình chuẩn bị dữ liệu (data_collator), mà trong hook chuyển tiếp của backbone. Điều này là do khi đi vào backbone, mã hóa đa phương thức từ phần ViT đã được hợp nhất với phần văn bản thuần túy, và embedding thu được tại thời điểm này có độ dài chính xác. Hơn nữa, thực hiện phân chia trong hook của backbone cũng tương thích với các mô hình văn bản thuần túy. Cách tiếp cận này cũng cho phép framework tránh lưu mã mô hình bổ sung, ngăn chặn chi phí bảo trì tăng lên khi mã mô hình gốc được cập nhật.

Thích ứng không đệm

Thích ứng không đệm có thể được hiểu là định dạng đầu vào ở dạng flash-attention: nhiều chuỗi được nối lại thành một siêu chuỗi dài.

Cách tiếp cận này mang lại rắc rối cho việc triển khai kỹ thuật thực tế. Do đó, trong quá trình triển khai, SWIFT áp dụng giải pháp kỹ thuật sau:

Phân rã đầu vào không đệm ban đầu, và thực hiện đệm (chia hết cho world_size*2) và phân chia riêng cho từng chuỗi.
Trước khi tính toán attention, theo vị trí đệm, đặt đệm của QV thành 0 và đệm của K thành các giá trị cực nhỏ để ngăn chặn đệm ảnh hưởng tiêu cực đến việc tính toán attention.
Do việc tính toán tổn thất cuối cùng của GRPO và DPO yêu cầu các chuỗi hoàn chỉnh, trong trường hợp không đệm, việc thu thập logits trước sẽ làm tăng khối lượng giao tiếp, trong khi thu thập sau sẽ gây ra tính toán tổn thất bất thường. Do đó, logic tính toán tổn thất cho từng phương pháp đào tạo cần được viết lại hoàn toàn.
Do Q chỉ có nửa độ dài khi số chuỗi giao tiếp lớn hơn thứ hạng, nó cần được khôi phục về độ dài đầy đủ trong quá trình cập nhật gradient ngược. Do đó, đệm grad cần được thực hiện riêng cho từng chuỗi, và LSE cần được đệm thành các giá trị cực nhỏ.

Lan truyền ngược

Theo quá trình suy luận công thức ở trên, quá trình lan truyền ngược của các cập nhật khối LSE và Attention-Out cần được thực hiện tuần tự và yêu cầu một số thông tin chuyển tiếp, chẳng hạn như LSE khối, Attention-Out khối, v.v. Mặc dù thông tin này có thể thu được trong quá trình chuyển tiếp flash_attn_forward, việc lưu trữ nó trong ctx có thể tiêu tốn thêm bộ nhớ. Do đó, chúng tôi chọn phương pháp tính toán lại flash_attn_forward một lần trong quá trình ngược, sau đó tính toán lse_grad dựa trên các kết quả trung gian, và sau đó thực hiện lan truyền ngược thực tế trên QKV.

Kết quả tối ưu hóa bộ nhớ

Chúng tôi đã sử dụng mô hình 3B và kiểm tra hiệu quả tối ưu hóa bộ nhớ trên 8 GPU A100:

bash NPROC_PER_NODE=8
swift sft
–model Qwen/Qwen2.5-3B-Instruct
–dataset ’test.jsonl’ \ # 9000 tokens per sequence –train_type lora
–torch_dtype bfloat16
–per_device_train_batch_size 4
–target_modules all-linear
–gradient_accumulation_steps 8
–save_total_limit 2
–save_only_model true
–save_steps 50
–max_length 65536
–warmup_ratio 0.05
–attn_impl flash_attn
–sequence_parallel_size 8
–logging_steps 1
–use_logits_to_keep false
–padding_free true

Như đã thấy ở phần đầu của bài viết, khi chia thành 8 phần, mức sử dụng bộ nhớ đào tạo giảm từ gần 80GiB xuống dưới 20GiB, đạt được hiệu quả mà các card đồ họa thương mại thông thường có thể thực hiện đào tạo.

Triển vọng

Bài viết này đã giải thích việc triển khai các khả năng đào tạo kết hợp Ulysses + Ring-Attention trong framework SWIFT. Chúng tôi vẫn đang tiếp tục khám phá các tối ưu hóa hơn nữa trong lĩnh vực này, ví dụ:

Trong quá trình ngược, liệu việc tính toán lại flash_attention_forward có phải là cách triển khai đạt được tốc độ tối ưu không?
Vẫn còn tiềm năng để tối ưu hóa hơn nữa về khối lượng giao tiếp P2P và các hướng thực thi không đồng bộ.

Các nhà phát triển quan tâm đến điều này có thể đưa ra các đề xuất có giá trị để giúp SWIFT cùng nhau cải thiện khả năng đào tạo mô hình lớn trong các tình huống chuỗi dài.

Tài liệu tham khảo:

Thảo luận

Locke

Could you provide some reference code? Using the trainer, I’m confused by the dataloader and DistributedSampler. Different ranks in the same sp_group always fail to obtain the same data idx es.

plkmn

USP also combine DeepSpeed-Ulysses and Ring-Attention together. https://arxiv.org/abs/2405.07719

Chú ý: Bài viết gốc được xuất bản vào ngày 16 tháng 9 năm 2025. Các nhận xét có thể không phản ánh thông tin mới nhất.

AI Today - SkyAI