하이브리드 샤딩된 데이터 병렬 처리

샤딩된 데이터 병렬 처리는 디바이스 간에 모델 상태(모델 파라미터, 그라데이션 및 옵티마이저 상태)를 분할하는 메모리 절약형 분산 훈련 기법입니다. 이렇게 하면 더 큰 모델에 맞추거나 여유 GPU 메모리를 사용하여 배치 크기를 늘릴 수 있습니다. SMP 라이브러리는 사용 GPUs 중인 전체 세트에 걸쳐 기본적으로 샤드를 PyTorch 완전히 샤드된 데이터 병렬(FSDP) PyTorch FSDP로 샤드된 데이터 병렬 처리를 실행하는 기능을 제공합니다. SMP v2에서 라이브러리는 , , FULL_SHARD, 에서 제공하는 샤딩 전략 중 하나인 PyTorch 하이브리드 샤딩(HYBRID_SHARD)을 확장하여 이 샤딩된 데이터 병렬 처리를 제공합니다 PyTorch FSDPSHARD_GRAD_OPHYBRID_SHARD_HYBRID_SHARD_ZERO2. PyTorch FSDP 이러한 방식으로 하이브리드 샤딩을 확장하면 블로그에 설명된 대로 에 대한 PyTorch gigantic-model 훈련의 거의 선형적인 조정을 AWS 구현 scale-aware-sharding하는 데 도움이 됩니다FSDP.

SMP 라이브러리를 사용하면 구성 가능한 수의 _HYBRID_SHARD_ZERO2에서 HYBRID_SHARD 및 를 쉽게 사용할 수 있으므로 단일 노드(HYBRID_SHARD) 또는 모든 GPUs ()에서 샤딩을 지원하는 네이티브 PyTorch FSDP를 GPUs확장할 수 있습니다FULL_SHARD. PyTorch FSDP 호출은 그대로 유지될 수 있으며 다음 코드 예제와 같이 SMP 구성에 hybrid_shard_degree 인수만 추가하면 됩니다. PyTorch 모델 주변의 FSDP 래퍼에서 sharding_strategy PyTorch 인수 값을 변경할 필요가 없습니다. ShardingStrategy.HYBRID_SHARD 를 값으로 전달할 수 있습니다. 또는 SMP 라이브러리는 스크립트의 전략을 재정의하고 hybrid_shard_degree 파라미터에 2보다 크거나 같은 값을 지정하는 ShardingStrategy.HYBRID_SHARD 경우 로 설정합니다.

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

SMP 구성 사전

{ "hybrid_shard_degree": 16 }

훈련 스크립트에서

import torch.sagemaker as tsm
tsm.init()

# Set up a PyTorch model
model = ...

# Wrap the PyTorch model using the PyTorch FSDP module
model = FSDP(
    model,
    ...
)

# Optimizer needs to be created after FSDP wrapper
optimizer = ...