활성화 오프로딩

일반적으로 순방향 패스는 각 계층의 활성화를 계산하고 해당 계층의 역방향 패스가 완료될 때까지 GPU 메모리에 유지합니다. 이러한 텐서를 순방향 패스 후 CPU 메모리로 오프로드하고 필요할 때 GPU로 다시 가져오면 상당한 GPU 메모리 사용량을 절약할 수 있습니다. PyTorch는 활성화 오프로드를 지원하지만, 구현을 통해 역방향 패스 중에 CPU에서 활성화를 가져오는 동안 GPU가 유휴 상태가 됩니다. 이로 인해 활성화 오프로드 사용 시 성능이 크게 저하됩니다.

SMP v2는 이 활성화 오프로드를 개선합니다. GPU가 해당 활성화에 대한 이전 패스를 시작하기 위해 활성화가 필요하기 전에 미리 활성화를 미리 가져옵니다. 사전 가져오기 기능을 사용하면 유휴 GPU 없이 훈련 진행 상황을 더욱 효율적으로 실행할 수 있습니다. 따라서 성능 저하 없이 메모리 사용량을 줄일 수 있습니다.

기본 PyTorch 모듈을 유지하여 훈련 스크립트에서 활성화를 오프로드할 수 있습니다. 다음은 스크립트에서 SMP 활성화 오프로드 기능을 적용하는 예제 구조입니다. 활성화 오프로드는 활성화 체크포인트과 함께 사용하는 경우에만 적용됩니다. 활성화 오프로드용 기본 PyTorch 체크포인트 도구에 대한 자세한 내용은 다음을 참조하세요.

PyTorch 활성화 체크포인트에 SMP 활성화 오프로드 기능을 적용할 수 있습니다. 이 작업은 2단계: 훈련 작업 시작 중에 sm_activation_offloading 및 activation_loading_horizon 파라미터를 SMP 구성 사전에 추가하여 수행됩니다.

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

SMP 구성

{
    "activation_loading_horizon": 2,
    "sm_activation_offloading": True
}

훈련 스크립트 내

import torch.sagemaker as tsm
tsm.init()

# Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing, 
    offload_wrapper,
)

model = FSDP(...)

# Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
    model,
    check_fn=checkpoint_transformer_layers_policy,
)

model = offload_wrapper(model)