Context Parallelism

컨텍스트 병렬화는 긴 입력 시퀀스를 여러 가속기에 분할하는 분산 컴퓨팅 기법입니다. 각 GPU는 시퀀스의 일부만 처리하면서 어텐션 연산 중 다른 GPU와 협력합니다. 이를 통해 장치별 활성화 메모리를 줄여 매우 긴 문서, 긴 비디오 또는 대규모 이미지 패치 집합과 같이 단일 GPU의 메모리 용량을 초과할 수 있는 입력으로 Transformer를 학습할 수 있습니다.

모델의 컨텍스트 윈도우를 단순히 늘리는 것과 달리, 컨텍스트 병렬화는 아키텍처가 이론적으로 받아들일 수 있는 정보량을 변경하지 않습니다. 대신 시퀀스 차원을 분산해 해당 컨텍스트를 계산상 효율적으로 처리할 수 있게 합니다.

컨텍스트 병렬화 작동 방식#

시퀀스가 32,000개의 토큰으로 구성되어 있고 컨텍스트 병렬화에 GPU 4개를 사용한다고 가정해 보겠습니다. 각 장치는 처음에 약 8,000개의 토큰을 받고 이에 해당하는 중간 활성화를 저장합니다.

정규화 및 피드포워드 레이어와 같은 대부분의 연산은 이러한 로컬 시퀀스 청크를 독립적으로 처리할 수 있습니다. 문제는 어텐션 메커니즘입니다. 로컬 쿼리는 다른 모든 장치에 저장된 키와 값에 어텐션해야 할 수 있습니다.

따라서 구현에서는 GPU 간에 키-값 또는 KV 블록을 교환합니다. 링 어텐션에서는 각 장치가 로컬 데이터를 사용해 부분 어텐션을 계산하고 KV 블록을 다음 장치로 전달한 뒤, 전체 시퀀스를 처리할 때까지 이 과정을 반복합니다. PyTorch 컨텍스트 병렬화 튜토리얼은 분산 스케일드 닷 프로덕트 어텐션을 통해 이 동작을 보여주며, NVIDIA 컨텍스트 병렬화 패키지는 올개더, 리듀스-스캐터 및 링 기반 통신 옵션을 설명합니다.

최종 결과는 일반적인 수치 차이를 제외하면 전체 시퀀스 어텐션과 수학적으로 동일하지만, 어느 단일 GPU도 모든 시퀀스 활성화를 보유할 필요가 없습니다.

컨텍스트 병렬화가 중요한 이유#

긴 시퀀스는 두 가지 주요 확장성 문제를 일으킵니다. 첫째, 시퀀스 길이가 길어질수록 저장된 활성화가 더 많은 메모리를 차지합니다. 둘째, 표준 셀프 어텐션은 시퀀스의 토큰을 서로 비교하므로 상당한 연산과 임시 데이터가 필요합니다.

컨텍스트 병렬화는 활성화를 여러 장치에 분할하여 메모리 문제를 해결합니다. 어텐션 연산도 분할할 수 있지만, 이 경우 통신 비용이 발생합니다. 효율적인 시스템은 NCCL 집합 연산 가이드에 문서화된 연산과 같이 KV 전송과 계산을 중첩합니다.

이 기법은 모델 가중치나 배치 크기가 아니라 시퀀스 길이 때문에 메모리 부족 오류가 발생할 때 가장 유용합니다. 컨텍스트 병렬화는 더 광범위한 분산 학습 분야에 속하며, 여러 차원을 동시에 확장하기 위해 다른 전략과 함께 사용하는 경우가 많습니다.

컨텍스트 병렬화와 관련 기법 비교#

• 텐서 병렬화는 개별 레이어 내의 연산이나 가중치 행렬을 나눕니다. 반면 컨텍스트 병렬화는 시퀀스 차원을 따라 토큰을 나눕니다.
• 파이프라인 병렬화는 모델 레이어의 여러 그룹을 서로 다른 장치에 할당합니다. 시퀀스 길이가 아니라 모델의 깊이를 분할합니다.
• 데이터 병렬화는 모델을 복제하고 각 복제본에 서로 다른 학습 예시를 제공합니다. PyTorch 분산 처리 개요는 전체 모델과 각 샘플이 하나의 GPU에 들어가는 경우 데이터 병렬화를 권장합니다.
• 시퀀스 병렬화는 텐서 병렬화와 관련된 일부 연산의 활성화를 분할하는 경우가 많습니다. 컨텍스트 병렬화는 네트워크 입력과 활성화 전반에 걸쳐 더 광범위하게 시퀀스를 분할합니다.

이러한 접근 방식은 서로 보완적입니다. NVIDIA 병렬화 전략 가이드는 컨텍스트, 텐서, 파이프라인 및 데이터 병렬화가 다차원 장치 배치를 구성하는 방법을 보여줍니다.

실제 적용 사례#

• 장문 문서 AI: 법률 또는 의료 언어 모델은 사건 기록 전체, 환자 이력 또는 기술 매뉴얼을 처리해야 할 수 있습니다. 컨텍스트 병렬화는 수천 개의 문서 토큰을 여러 가속기에 분산해 섹션 간 어텐션을 유지하면서 활성화 메모리 부담을 줄입니다.

• 장시간 비디오 및 멀티모달 이해: 비디오 Transformer와 대규모 비전 모델은 프레임, 이미지 패치, 오디오 세그먼트 및 텍스트를 하나의 긴 토큰 시퀀스로 표현할 수 있습니다. 이 시퀀스를 분산하면 프레임 수나 공간적 세부 정보를 크게 줄이지 않고도 모델이 긴 녹화물을 분석할 수 있습니다. AWS Neuron 컨텍스트 병렬화 개요는 이러한 긴 컨텍스트 워크로드에서 가속기 그룹이 KV 샤드를 교환하는 방식을 보여줍니다.

Ultralytics YOLO26과 같은 소형 컴퓨터 비전 아키텍처에는 일반적으로 컨텍스트 병렬화가 필요하지 않습니다. 대개 Ultralytics 모델 학습 워크플로를 통한 표준 멀티 GPU 데이터 병렬화가 학습 가속에 더 적절합니다.

실용적인 사용과 절충점#

다음의 최소 예시는 PyTorch의 실험적 컨텍스트 병렬화 API와 스케일드 닷 프로덕트 어텐션 제어 기능을 사용합니다. 이 예시를 cp_example.py로 저장한 다음 torchrun --standalone --nproc-per-node=2 cp_example.py를 사용해 GPU 2개에서 실행합니다.

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

여기서 차원 2는 시퀀스 차원이므로 각 프로세스는 시퀀스 샤드를 받고, 어텐션은 장치 메시 전반에 걸쳐 조정됩니다.

실제로 엔지니어는 감소된 메모리 사용량이 통신 오버헤드보다 큰지 확인하고, 고속 인터커넥트를 사용하며, 대표적인 시퀀스 길이를 벤치마크해야 합니다. 표준 비전 프로젝트의 경우 Ultralytics Platform은 수동 컨텍스트 병렬화 구성 없이 데이터셋 어노테이션, 학습, 배포 및 모니터링을 위한 더 간단한 클라우드 및 로컬 워크플로를 제공합니다.