활성화 오프로딩

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

SMP v2는 이 활성화 오프로드를 개선합니다. 가 해당 활성화에 대해 역방향 패스를 시작하는 GPU 데 활성화가 필요하기 전에 미리 활성화를 미리 가져옵니다. 미리 가져오기 기능은 유휴 상태 없이 훈련 진행 상황을 더 효율적으로 실행하는 데 도움이 됩니다GPUs. 이로 인해 성능 저하 없이 메모리 사용량을 줄일 수 있습니다.

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

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

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

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)