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

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

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

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

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 = ...