Context Parallelism

Context parallelism là một kỹ thuật tính toán phân tán giúp chia một chuỗi đầu vào dài ra trên nhiều bộ gia tốc. Mỗi GPU chỉ xử lý một phần của chuỗi trong khi phối hợp với các GPU khác trong quá trình tính toán attention. Điều này giúp giảm bộ nhớ activation trên mỗi thiết bị, cho phép một transformer huấn luyện trên các đầu vào vượt quá dung lượng bộ nhớ của một GPU, chẳng hạn như các tài liệu rất dài, video kéo dài, hoặc tập hợp lớn các bản vá hình ảnh.

Không giống như việc chỉ đơn giản mở rộng context window của một model, context parallelism không làm thay đổi lượng thông tin mà kiến trúc có thể chấp nhận về mặt lý thuyết. Thay vào đó, nó làm cho việc xử lý ngữ cảnh đó trở nên khả thi về mặt tính toán bằng cách phân phối chiều chuỗi (sequence dimension).

Cách thức hoạt động của Context Parallelism#

Giả sử một chuỗi chứa 32.000 token và context parallelism sử dụng bốn GPU. Ban đầu, mỗi thiết bị nhận khoảng 8.000 token và lưu trữ các activation trung gian tương ứng.

Hầu hết các phép toán, chẳng hạn như normalization và các feed-forward layer, có thể xử lý các đoạn chuỗi cục bộ này một cách độc lập. Thách thức nằm ở attention mechanism: một query cục bộ có thể cần attend tới các key và value do mọi thiết bị khác nắm giữ.

Do đó, các quá trình triển khai sẽ trao đổi các khối key-value (hay KV) giữa các GPU. Trong ring attention, mỗi thiết bị tính toán partial attention bằng dữ liệu cục bộ của mình, chuyển một khối KV sang thiết bị tiếp theo, và lặp lại cho đến khi xử lý xong toàn bộ chuỗi. PyTorch context parallel tutorial minh họa hành vi này thông qua distributed scaled dot-product attention, trong khi NVIDIA context parallel package mô tả các tùy chọn giao tiếp all-gather, reduce-scatter và ring-based.

Kết quả cuối cùng tương đương về mặt toán học với full-sequence attention, chịu ảnh hưởng bởi các sai lệch số học thông thường, nhưng không một GPU đơn lẻ nào phải giữ lại toàn bộ activation của chuỗi.

Tại sao Context Parallelism lại quan trọng#

Các chuỗi dài tạo ra hai vấn đề mở rộng lớn. Đầu tiên, các activation được lưu trữ tiêu tốn nhiều bộ nhớ hơn khi chiều dài chuỗi tăng lên. Thứ hai, self-attention chuẩn so sánh các token xuyên suốt chuỗi, tạo ra lượng tính toán lớn và dữ liệu tạm thời.

Context parallelism giải quyết vấn đề bộ nhớ bằng cách chia nhỏ các activation giữa các thiết bị. Nó cũng có thể phân chia công việc attention, mặc dù quá trình giao tiếp tạo ra một chi phí mới. Các hệ thống hiệu quả sẽ chồng chéo việc truyền KV với quá trình tính toán bằng cách sử dụng các thao tác như được tài liệu hóa trong NCCL collective operations guide.

Kỹ thuật này có giá trị nhất khi chiều dài chuỗi, chứ không phải trọng số của model hay kích thước batch, gây ra lỗi out-of-memory. Nó thuộc lĩnh vực rộng lớn hơn về distributed training và thường được kết hợp với các chiến lược khác để mở rộng quy mô đa chiều cùng lúc.

Context Parallelism so với các kỹ thuật liên quan#

• Tensor parallelism chia các phép toán hoặc ma trận trọng số trong các layer riêng lẻ. Ngược lại, context parallelism chia các token dọc theo chiều chuỗi.
• Pipeline parallelism gán các nhóm layer khác nhau của model cho các thiết bị khác nhau. Nó phân vùng chiều sâu của model thay vì chiều dài chuỗi.
• Data parallelism nhân bản model và cung cấp cho mỗi bản sao các ví dụ huấn luyện khác nhau. PyTorch distributed overview đề xuất phương pháp này khi toàn bộ model và mỗi mẫu dữ liệu đều vừa vặn trên một GPU.
• Sequence parallelism thường chia nhỏ (shard) các activation cho các phép toán được chọn liên quan đến tensor parallelism. Context parallelism áp dụng việc phân vùng chuỗi rộng rãi hơn trên các đầu vào và activation của mạng.

Các phương pháp này mang tính bổ trợ cho nhau. NVIDIA parallelism strategies guide cho thấy cách context, tensor, pipeline và data parallelism có thể tạo thành một bố cục thiết bị đa chiều.

Các ứng dụng trong thực tế#

• AI tài liệu dài: Một ngôn ngữ model về pháp lý hoặc y tế có thể cần xử lý toàn bộ hồ sơ vụ án, tiền sử bệnh án hoặc tài liệu kỹ thuật. Context parallelism phân phối hàng nghìn token tài liệu trên các bộ gia tốc, giảm áp lực bộ nhớ activation trong khi vẫn duy trì attention giữa các phần ở xa nhau.

• Hiểu video dài và đa phương thức: Các video transformer và large vision models có thể biểu diễn các khung hình, bản vá hình ảnh, đoạn âm thanh và văn bản thành một chuỗi token dài. Việc phân phối chuỗi đó giúp model phân tích các bản ghi kéo dài mà không cần giảm quá mạnh số lượng khung hình hoặc chi tiết không gian. AWS Neuron context parallelism overview minh họa cách các nhóm bộ gia tốc có thể trao đổi các shard KV cho các khối lượng công việc ngữ cảnh dài này.

Đối với các kiến trúc thị giác máy tính nhỏ gọn như Ultralytics YOLO26, context parallelism thường là không cần thiết. Data parallelism đa GPU tiêu chuẩn thông qua Ultralytics model training workflow thường là cách thích hợp hơn để tăng tốc huấn luyện.

Sử dụng thực tế và các đánh đổi#

Ví dụ tối giản sau đây sử dụng API context-parallel thử nghiệm của PyTorch và scaled dot-product attention controls. Lưu nó dưới dạng cp_example.py và chạy trên hai GPU với torchrun --standalone --nproc-per-node=2 cp_example.py.

import os

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.experimental import context_parallel
from torch.nn.attention import SDPBackend, sdpa_kernel

rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])

torch.cuda.set_device(rank)
torch.cuda.manual_seed(0)
dist.init_process_group("nccl")

mesh = init_device_mesh("cuda", (world_size,))
qkv = [torch.randn(1, 4, 4096, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) for _ in range(3)]

with sdpa_kernel(SDPBackend.FLASH_ATTENTION), context_parallel(mesh, buffers=tuple(qkv), buffer_seq_dims=(2, 2, 2)):
    output = F.scaled_dot_product_attention(*qkv, is_causal=True)

output.float().square().mean().backward()
dist.destroy_process_group()

Ở đây, chiều 2 là chiều chuỗi, do đó mỗi tiến trình nhận một shard chuỗi trong khi attention phối hợp xuyên suốt lưới thiết bị.

Trong thực tế, các kỹ sư nên xác nhận rằng việc giảm bộ nhớ lớn hơn chi phí giao tiếp, sử dụng các kết nối liên mạng tốc độ cao, và đánh giá điểm chuẩn (benchmark) các độ dài chuỗi tiêu biểu. Đối với các dự án thị giác chuẩn, Ultralytics Platform cung cấp các quy trình làm việc trên đám mây và cục bộ đơn giản hơn cho việc gán nhãn tập dữ liệu, huấn luyện, triển khai và giám sát mà không yêu cầu cấu hình context-parallel thủ công.