컨텍스트 병렬 처리

컨텍스트 병렬 처리는 시퀀스 차원을 따라 모델 활성화를 분할하는 모델 병렬 처리의 한 유형입니다. LayerNorm 및 RMSNorm만 분할하는 다른 시퀀스 병렬 처리 기술과 달리 컨텍스트 병렬 처리는 시퀀스 차원을 따라 네트워크 입력 및 모든 중간 활성화를 분할합니다.

SMP v2는 컨텍스트 병렬 처리를 위해 Transformer Engine과 통합되며 PyTorch FSDP 및 SMP 텐서 병렬화와 함께 사용할 수 있습니다. 모델 훈련을 위해 세 개의 병렬 처리를 모두 동시에 활성화할 수 있습니다. 컨텍스트 병렬 처리는 활성화 크기가 크고 시퀀스 길이가 긴 훈련 모델에 유용합니다. 각 디바이스가 시퀀스 차원을 따라 점수 및 출력의 일부만 계산할 수 있도록 하여 주의 점수 및 주의 출력의 계산을 가속화합니다. 또한 텐서 병렬 처리는 숨겨진 차원을 따라 분할을 통해 계산을 가속화하지만, 컴퓨팅 요구 사항이 시퀀스 차원에 따라 사분면으로 증가하므로 컨텍스트 병렬 처리의 이점이 더 큽니다.

SMP 컨텍스트 병렬 처리와 호환되는 Hugging Face 트랜스포머 모델

SMP v2는 현재 다음 Hugging Face 트랜스포머 모델에 대한 컨텍스트 병렬 처리를 지원합니다.

• GPT-NeoX

• Llama 2 및 Llama 3

• Mistral 7B

컨텍스트 병렬 처리 구성

클러스터의 GPU 수를 균등하게 나눈 context_parallel_degree 파라미터로 정수 값을 설정합니다. 예를 들어 8-GPU 인스턴스가 있는 경우 context_parallel_degree에 2, 4 또는 8을 사용합니다. 작은 context_parallel_degree 값으로 시작하여 모델이 필요한 입력 시퀀스 길이로 GPU 메모리에 적합할 때까지 점진적으로 늘리는 것이 좋습니다.

다음 코드 조각은 SageMaker 모델 병렬 처리 라이브러리 v2 사용에 도입된 2단계 프로세스를 따르면서 훈련 스크립트에 SMP 초기화 모듈 torch.sagemaker.init()를 추가하고 훈련 작업 시작 관리자를 위한 JSON 형식의 SMP 구성 사전을 설정하는 방법을 보여줍니다. PyTorch 모델 또는 PyTorch FSDP 구성을 변경할 필요가 없습니다. context_parallel_degree 파라미터에 대한 자세한 내용은 SMP v2 코어 기능 구성 파라미터를 참조하세요.

훈련 스크립트에서

1단계의 일환으로 스크립트를 로 초기화torch.sagemaker.init()하여 SMP v2를 활성화하고 torch.sagemaker.transform API로 모델을 래핑합니다.

SMP v2.6.0부터 cp_comm_type 인수를 사용하여 사용할 컨텍스트 병렬 처리 구현을 결정할 수 있습니다. SMP 라이브러리는 현재 p2p 및 all_gather의 두 가지 구현을 지원합니다. 이 p2p 구현은 주의 구현 중에 키-값 누적에 대해 피어-투-피어 전송-수신 호출을 사용하고 비동기적으로 실행되어 컴퓨팅과 중복됩니다. all_gather 대신 구현은 AllGather 집합 작업을 사용하고 동기적으로 실행됩니다.

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model, cp_comm_type="p2p")

SMP 구성

2단계의 일부로 SageMaker PyTorch 추정기의 SMP 구성 사전에 다음 파라미터를 추가합니다.

{   
    ..., # other SMP config parameters
    "context_parallel_degree": 2
}