Ring Attention

Ring attention tính toán self-attention chính xác trên một chuỗi quá dài đối với một accelerator bằng cách chia chuỗi thành các block trải trên những thiết bị được sắp xếp thành vòng. Mỗi thiết bị giữ một block truy vấn và chuyển các block khóa cùng giá trị cho thiết bị kế tiếp, đồng thời chồng lấp hoạt động truyền dữ liệu với phép tính trên block đang giữ; nhờ đó, độ dài chuỗi tối đa mà model xử lý được tăng tỷ lệ thuận với số thiết bị. Hao Liu, Matei Zaharia và Pieter Abbeel tại UC Berkeley giới thiệu phương pháp này vào tháng 10 năm 2023 trong bài báo Ring Attention with Blockwise Transformers for Near-Infinite Context, phát triển dựa trên transformer song song theo block trước đó của họ; các tác giả báo cáo việc huấn luyện và suy luận với cửa sổ ngữ cảnh hàng triệu token mà không xấp xỉ attention.

Cách Ring Attention hoạt động#

Chuỗi được chia thành N block, mỗi block nằm trên một GPU hoặc TPU. Mỗi thiết bị giữ block truy vấn của mình trong toàn bộ phép tính. Qua N bước, thiết bị tính attention giữa truy vấn của nó và block khóa-giá trị hiện đang nắm giữ, gửi block đó đến thiết bị kế tiếp và nhận một block khác từ thiết bị trước đó:

Sau N bước, mỗi block truy vấn đã lần lượt kết hợp với mọi block khóa và giá trị. Không thiết bị nào phải lưu toàn bộ ma trận attention hay toàn bộ tập khóa và giá trị, vì vậy bộ nhớ trên mỗi thiết bị phụ thuộc vào kích thước block thay vì độ dài chuỗi. Hoạt động truyền dữ liệu được che khuất bởi phép tính miễn là tính attention trên một block mất nhiều thời gian hơn truyền block đó, từ đó xác định kích thước block tối thiểu cho phần cứng cụ thể. Các kết quả trung gian được kết hợp bằng cùng phép softmax tích lũy mà flash attention sử dụng bên trong một GPU: mỗi thiết bị giữ tổng tích lũy của các điểm số đã lấy hàm mũ và các giá trị có trọng số, rồi cuối cùng lấy tổng thứ nhất chia cho tổng thứ hai. Mô phỏng bốn thiết bị trên một máy cho thấy kết quả chính xác:

import torch

torch.manual_seed(0)
n_devices, block, dim = 4, 8, 16
q, k, v = torch.randn(3, n_devices * block, dim)
q_blocks, k_blocks, v_blocks = q.chunk(n_devices), k.chunk(n_devices), v.chunk(n_devices)

outputs = []
for i in range(n_devices):  # device i keeps query block i for the whole computation
    numerator, denominator = torch.zeros(block, dim), torch.zeros(block, 1)
    for step in range(n_devices):
        j = (i - step) % n_devices  # key-value block that has arrived from the previous device
        p = torch.exp(q_blocks[i] @ k_blocks[j].T / dim**0.5)
        numerator += p @ v_blocks[j]
        denominator += p.sum(dim=-1, keepdim=True)
    outputs.append(numerator / denominator)

full = torch.softmax(q @ k.T / dim**0.5, dim=-1) @ v
print(f"Largest score block per device: {block}x{block}, full matrix: {len(q)}x{len(k)}")
print("Matches full attention:", torch.allclose(torch.cat(outputs), full, atol=1e-5))

Kết quả in ra Matches full attention: True: mỗi thiết bị chỉ giữ một block điểm số 8 × 8, và khi kết hợp, chúng tạo attention trên toàn bộ ma trận 32 × 32. Các kernel production cũng trừ đi giá trị cực đại tích lũy trước khi lấy hàm mũ để giữ cho các con số ổn định. Với causal mask, các thiết bị giữ những block truy vấn đầu chuỗi bỏ qua hầu hết block khóa và phải chờ; Striped Attention khắc phục điều này bằng cách phân bổ cho mỗi thiết bị các token trải đều trong chuỗi, tăng thông lượng huấn luyện đầu-cuối lên đến 1.45 lần ở 256K token.

Ứng dụng#

Ring attention hữu ích ở bất cứ đâu có đầu vào quá dài đối với một thiết bị đơn lẻ, và video là trường hợp rõ ràng nhất. Cùng nhóm Berkeley đã dùng phương pháp này để huấn luyện Large World Model (LWM), một họ mô hình thị giác-ngôn ngữ (VLM) 7 tỷ tham số, với ngữ cảnh tăng dần từ 4K lên 1 triệu token văn bản và video xen kẽ. LWM mã hóa mỗi khung hình thành 256 token, vì vậy một giờ video lấy mẫu một khung hình mỗi giây tương ứng 921,600 token; model đã trả lời câu hỏi về các video tổng hợp dài một giờ trên YouTube, một tác vụ hiểu video cần nắm bắt ngữ cảnh tầm xa. Việc huấn luyện mô hình ngôn ngữ lớn (LLM) đọc toàn bộ codebase hoặc bộ sưu tập tài liệu trong một prompt cũng gặp cùng giới hạn. API context_parallel() của PyTorch thay attention tích vô hướng có tỷ lệ bằng ring attention trong context của nó; nhóm TorchTitan đã dùng phương pháp này để huấn luyện Llama 3 8B trên các chuỗi 1 triệu token qua 32 GPU H100.

Ring Attention và các kiểu song song hóa khác#

• Song song hóa ngữ cảnh: Chiến lược tổng quát để chia một chuỗi qua nhiều thiết bị. Ring attention là thuật toán đứng sau triển khai của PyTorch; DeepSpeed Ulysses là phương án thay thế, sử dụng trao đổi all-to-all để cung cấp cho mỗi thiết bị toàn bộ token của một tập con các attention head.
• Song song hóa tensor và song song hóa pipeline: Chia trọng số model bên trong các layer hoặc giữa các layer, hữu ích khi số tham số vượt quá bộ nhớ của một thiết bị.
• Flash attention: Kernel chạy trên một thiết bị. Ring attention chạy một kernel xử lý theo block thuộc loại này trên mỗi thiết bị cho block cục bộ.

Song song hóa dữ liệu chia batch thay vì chia chuỗi, và đây là chế độ mà Ultralytics YOLO26 sử dụng để huấn luyện đa GPU: truyền device=[0, 1] sẽ khởi chạy PyTorch DistributedDataParallel, duy trì một bản sao model đầy đủ trên mỗi GPU, phân bổ cho mỗi GPU một phần bằng nhau của từng batch ảnh và lấy trung bình gradient sau mỗi bước.