활성화 오프로딩

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

SMPv2는 이러한 활성화 오프로딩을 개선합니다. 활성화에 대한 역방향 전달을 시작하기 위해 활성화가 필요하기 전에 미리 활성화를 미리 가져옵니다. GPU 프리페칭 기능을 사용하면 유휴 상태 없이 훈련 진행 상황을 더 효율적으로 실행할 수 있습니다. GPUs 따라서 성능 저하 없이 메모리 사용량을 줄이는 이점을 얻을 수 있습니다.

활성화 오프로드를 위한 네이티브 PyTorch 모듈을 교육 스크립트에 그대로 둘 수 있습니다. 다음은 스크립트에 SMP 활성화 오프로드 기능을 적용하는 예제 구조입니다. 참고로 활성화 오프로드는 와 함께 사용하는 경우에만 적용됩니다. 활성화 체크포인트 활성화 오프로딩을 위한 기본 PyTorch 체크포인트 도구에 대한 자세한 내용은 다음을 참조하십시오.

활성화 체크포인트에 SMP 활성화 오프로드 기능을 적용할 수 있습니다. PyTorch 이 작업을 수행하는 동안 SMP 구성 사전에 sm_activation_offloading 및 activation_loading_horizon 매개변수를 추가하면 됩니다. 2단계: 교육 작업 시작

다음 코드 스니펫은 에서 소개한 2단계 프로세스를 따르면서 교육 스크립트에 SMP 초기화 모듈을 torch.sagemaker.init() 추가하고 교육 작업 시작 프로그램을 위한 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)