Context Parallelism

컨텍스트 병렬성은 긴 입력 시퀀스를 여러 가속기에 분산시키는 분산 컴퓨팅 기술입니다. 각 GPU는 어텐션 연산 중에 다른 GPU와 협력하면서 시퀀스의 일부만 처리합니다. 이를 통해 디바이스당 활성화 메모리가 줄어들어, 매우 긴 문서, 확장된 비디오 또는 대규모 이미지 패치 컬렉션과 같이 단일 GPU의 메모리를 초과할 수 있는 입력에 대해 transformer를 학습시킬 수 있습니다.

모델의 context window를 단순히 확장하는 것과 달리, 컨텍스트 병렬성은 아키텍처가 이론적으로 수용할 수 있는 정보의 양을 변경하지 않습니다. 대신 시퀀스 차원을 분산시켜 해당 컨텍스트를 처리하는 것을 계산상 실현 가능하게 만듭니다.

컨텍스트 병렬성의 작동 방식#

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

정규화 및 피드포워드 레이어와 같은 대부분의 연산은 이러한 로컬 시퀀스 청크를 독립적으로 처리할 수 있습니다. 문제는 attention mechanism입니다. 로컬 쿼리는 다른 모든 디바이스가 보유한 키와 값을 참조해야 할 수 있습니다.

따라서 구현체에서는 GPU 간에 키-값(KV) 블록을 교환합니다. ring attention에서는 각 디바이스가 로컬 데이터를 사용하여 부분 어텐션을 계산하고, 다음 디바이스로 KV 블록을 전달하며, 전체 시퀀스를 처리할 때까지 이를 반복합니다. PyTorch context parallel tutorial은 분산된 스케일드 닷 프로덕트 어텐션을 통해 이 동작을 보여주며, NVIDIA context parallel package는 all-gather, reduce-scatter 및 링 기반 통신 옵션을 설명합니다.

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

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

긴 시퀀스는 두 가지 주요 스케일링 문제를 야기합니다. 첫째, 시퀀스 길이가 길어짐에 따라 저장된 활성화 값이 더 많은 메모리를 소비합니다. 둘째, 표준 셀프 어텐션은 시퀀스 전체에서 토큰을 비교하므로 상당한 계산량과 임시 데이터가 발생합니다.

컨텍스트 병렬성은 디바이스 간에 활성화 값을 분할하여 메모리 문제를 해결합니다. 또한 통신 비용이 추가되지만 어텐션 작업을 분할할 수도 있습니다. 효율적인 시스템은 NCCL collective operations guide에 문서화된 작업과 같은 연산을 사용하여 KV 전송과 계산을 중첩합니다.

이 기술은 모델 가중치나 배치 크기 대신 시퀀스 길이로 인해 OOM(Out-Of-Memory) 오류가 발생할 때 가장 유용합니다. 이는 distributed training의 광범위한 분야에 속하며, 여러 차원을 동시에 스케일링하기 위해 다른 전략과 결합되는 경우가 많습니다.

컨텍스트 병렬성 대 관련 기술#

• **Tensor parallelism**은 개별 레이어 내의 연산이나 가중치 행렬을 분할합니다. 반면 컨텍스트 병렬성은 시퀀스 차원을 따라 토큰을 분할합니다.
• **Pipeline parallelism**은 모델 레이어의 서로 다른 그룹을 서로 다른 디바이스에 할당합니다. 시퀀스 길이가 아니라 모델 깊이를 분할합니다.
• 데이터 병렬성은 모델을 복제하고 각 복제본에 서로 다른 학습 예제를 제공합니다. PyTorch distributed overview에서는 전체 모델과 각 샘플이 하나의 GPU에 들어갈 때 이를 권장합니다.
• 시퀀스 병렬성은 종종 텐서 병렬성과 관련된 선택된 연산에 대한 활성화 값을 샤딩합니다. 컨텍스트 병렬성은 네트워크 입력 및 활성화 전반에 걸쳐 시퀀스 분할을 더 광범위하게 적용합니다.

이러한 접근 방식들은 상호 보완적입니다. NVIDIA parallelism strategies guide는 컨텍스트, 텐서, 파이프라인 및 데이터 병렬성이 다차원 디바이스 레이아웃을 형성할 수 있는 방법을 보여줍니다.

실제 애플리케이션 사례#

• 긴 문서 AI: 법률 또는 의료 언어 모델은 전체 케이스 파일, 환자 병력 또는 기술 설명서를 처리해야 할 수 있습니다. 컨텍스트 병렬성은 수천 개의 문서 토큰을 가속기에 분산시켜, 멀리 떨어진 섹션 간의 어텐션을 유지하면서 활성화 메모리 압박을 줄여줍니다.

• 긴 비디오 및 멀티모달 이해: 비디오 transformer 및 large vision models은 프레임, 이미지 패치, 오디오 세그먼트, 텍스트를 하나의 긴 토큰 시퀀스로 표현할 수 있습니다. 해당 시퀀스를 분산시키면 프레임 수나 공간적 디테일을 과도하게 줄이지 않고도 모델이 확장된 레코딩을 분석하는 데 도움이 됩니다. AWS Neuron context parallelism overview는 이러한 긴 컨텍스트 워크로드를 위해 가속기 그룹이 KV 샤드를 교환할 수 있는 방법을 보여줍니다.

Ultralytics YOLO26과 같은 컴팩트한 컴퓨터 비전 아키텍처의 경우 컨텍스트 병렬성이 일반적으로 불필요합니다. Ultralytics model training workflow를 통한 표준 멀티 GPU 데이터 병렬성이 일반적으로 학습을 가속화하는 더 적절한 방법입니다.

실제 사용 및 트레이드오프#

다음의 최소한의 예제는 PyTorch의 실험적인 컨텍스트 병렬 API와 scaled dot-product attention controls을 사용합니다. cp_example.py로 저장하고 torchrun --standalone --nproc-per-node=2 cp_example.py를 사용하여 두 개의 GPU에서 실행합니다.

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은 수동 컨텍스트 병렬 구성 없이도 데이터셋 주석, 학습, 배포 및 모니터링을 위한 더 간단한 클라우드 및 로컬 워크로드를 제공합니다.