SageMaker 분산 모델 병렬화 라이브러리 구성 팁과 함정

라이브러리는 배치 크기가 커질 때 가장 효율적입니다. 모델이 단일 기기에 적합하지만 배치 크기가 작아야만 훈련할 수 있는 사용 사례의 경우 라이브러리를 통합한 후에 배치 크기를 늘릴 수 있으며 그렇게 해야 합니다. 모델 병렬화를 통해 대형 모델의 경우 메모리가 절약되므로 이전에는 메모리에 담을 수 없었던 배치 크기를 사용하여 훈련할 수 있습니다.

너무 작거나 너무 큰 마이크로 배치를 여러 개 선택하면 성능이 저하될 수 있습니다. 라이브러리는 각 장치에서 각 마이크로 배치를 순차적으로 실행하므로 각 GPU를 충분히 활용할 수 있을 만큼 마이크로 배치 크기(배치 크기를 마이크로 배치 수로 나눈 값)가 충분히 커야 합니다. 동시에 마이크로 배치 수에 따라 파이프라인 효율성이 증가하므로 적절한 균형을 유지하는 것이 중요합니다. 일반적으로 2~4개의 마이크로 배치를 시도하여 배치 크기를 메모리 한도까지 늘린 다음 더 큰 배치 크기와 마이크로 배치 수를 실험해 보는 것이 좋습니다. 마이크로 배치 수가 증가하면 인터리브 파이프라인을 사용하면 배치 크기를 늘리는 것도 가능해질 수 있습니다.

배치 크기는 항상 마이크로 배치 수로 나눌 수 있어야 합니다. 데이터셋의 크기에 따라 때때로 모든 에포크의 마지막 배치의 크기가 나머지 배치보다 작을 수 있으며, 이 작은 배치도 마이크로 배치 수로 나눌 수 있어야 합니다. 그렇지 않은 경우 tf.Dataset.batch() 호출에서 설정 (drop_remainder=Truein) 하거나 set drop_last=True in (in TensorFlow) 으로 설정하여 이 마지막 소량 배치를 사용하지 않도록 할 수 있습니다. DataLoader PyTorch 데이터 파이프라인에 다른 API를 사용하는 경우 마지막 배치를 마이크로 배치 수로 나눌 수 없을 때마다 수동으로 건너뛰어야 할 수 있습니다.

수동 파티셔닝을 사용하는 경우 변환기 아키텍처의 임베딩 테이블과 같이 모델의 여러 작업 및 모듈에서 사용되는 파라미터에 유의하세요. 동일한 파라미터를 공유하는 모듈은 정확성을 위해 동일한 기기에 배치해야 합니다. 자동 파티셔닝을 사용하는 경우 라이브러리는 이 제약 조건을 자동으로 적용합니다.

모델이 여러 입력을 받는 경우 smp.dp_rank()로 데이터 파이프라인의 임의 연산(예: 셔플링)을 사용하여 시드해야 합니다. 데이터 세트를 여러 데이터 병렬 기기에서 결정론적으로 샤드하는 경우 샤드가 smp.dp_rank()에 의해 인덱싱되었는지 확인하세요. 이는 모델 파티션을 구성하는 모든 순위에서 표시되는 데이터 순서가 일관되도록 하기 위한 것입니다.

smp.DistributedModel.call(for TensorFlow) 또는 smp.DistributedModel.forward (for PyTorch) 함수에서 반환되는 모든 텐서는 해당 특정 텐서를 계산한 랭크에서 다른 모든 랭크로 브로드캐스트됩니다. 따라서 호출 및 전달 메서드 외부에서 필요하지 않은 텐서(예: 중간 활성화)는 반환되지 않아야 합니다. 이렇게 하면 불필요한 통신 및 메모리 오버헤드가 발생하고 성능이 저하되기 때문입니다.

smp.step으로 데코레이션된 함수에 일괄 처리 차원이 없는 텐서 인수가 있는 경우 smp.step 호출 시 인수 이름을 non_split_inputs 목록에 제공해야 합니다. 이렇게 하면 라이브러리가 텐서를 마이크로 배치로 분할하려고 시도하지 못하게 됩니다. 자세한 내용은 API 설명서의 smp.step를 참조하세요.

결정적 결과를 얻기 위해 시드를 사용하는 경우 smp.dp_rank()(예: torch.manual_seed(42 + smp.dp_rank()))를 기반으로 시드를 설정하세요. 이렇게 하지 않으면 nn.Parameter의 여러 파티션이 같은 방식으로 초기화되므로 컨버전스에 영향을 줍니다.

SageMaker의 모델 병렬화 라이브러리는 NCCL을 사용하여 모듈 배포에 필요한 집합체를 구현합니다. 특히 소형 모델의 경우 GPU에 너무 많은 NCCL 호출이 동시에 예약되면 NCCL에서 사용하는 추가 공간으로 인해 메모리 사용량이 증가할 수 있습니다. 이 문제를 해결하기 위해 특정 시간에 진행 중인 NCCL 작업 수가 지정된 smp 한도보다 작거나 같도록 NCCL 호출을 조절합니다. 기본 제한은 8이지만 환경 변수 SMP_NCCL_THROTTLE_LIMIT를 사용하여 조정할 수 있습니다. 텐서 병렬화를 사용하는 동안 메모리 사용량이 예상보다 많으면 이 제한을 줄여 볼 수 있습니다. 하지만 제한을 너무 작게 선택하면 처리량 손실이 발생할 수 있습니다. 전송률 조절을 완전히 비활성화하려면 SMP_NCCL_THROTTLE_LIMIT=-1로 설정할 수 있습니다.

텐서 병렬도가 1일 때 유지되는 다음 자격 증명은 텐서 병렬도가 1보다 크면 유지되지 않습니다. smp.mp_size() * smp.dp_size() == smp.size(). 이는 텐서 병렬 그룹이 모델 병렬화 그룹과 데이터 병렬화 그룹 모두에 속하기 때문입니다. 코드에 mp_rank, mp_size, MP_GROUP 등에 대한 기존 참조가 있고 파이프라인 병렬 그룹으로만 작업하려는 경우 참조를 smp.pp_size()로 바꿔야 할 수 있습니다. 다음 ID는 항상 유효합니다.

텐서 병렬화가 활성화되면 smp.DistributedModel 래퍼가 모델 파라미터를 수정하므로 smp.DistributedModel 호출 후 분산 파라미터를 사용하여 옵티마이저를 생성해야 합니다. 예를 들어 IAM은 다음을 허용하지 않습니다.

## WRONG
model = MyModel()
optimizer = SomeOptimizer(model.parameters())
model = smp.DistributedModel(model)  # optimizer now has outdated parameters! 

대신 다음과 같이 smp.DistributedModel 파라미터를 사용하여 옵티마이저를 생성해야 합니다.

## CORRECT
model = smp.DistributedModel(MyModel())
optimizer = SomeOptimizer(model.optimizers())

텐서 병렬화를 통해 모듈이 분산된 모듈로 대체되는 경우 분산 모듈은 원래 모듈에서 가중치를 상속하지 않고 새 가중치를 초기화합니다. 즉, 예를 들어 특정 호출(예: load_state_dict 호출을 통해)에서 가중치를 초기화해야 하는 경우 smp.DistributedModel 호출 후, 즉 모듈 배포가 이루어진 후에 초기화해야 합니다.

분산 모듈의 파라미터에 직접 액세스할 때는 가중치의 모양이 원래 모듈과 같지 않다는 점에 유의하세요. 예를 들면 다음과 같습니다. 

with smp.tensor_parallelism():
    linear = nn.Linear(60, 60)

# will pass
assert tuple(linear.weight.shape) == (60, 60)

distributed_linear = smp.DistributedModel(linear)

# will fail. the number of input channels will have been divided by smp.tp_size()
assert tuple(distributed_linear.module.weight.shape) == (60, 60)

텐서 병렬화에는 torch.utils.data.distributed.DistributedSampler 사용이 강력히 권장됩니다. 이렇게 하면 모든 데이터 병렬 순위에서 동일한 수의 데이터 샘플을 수신할 수 있으므로 여러 단계에서 다른 dp_rank를 수행할 때 발생할 수 있는 중단이 방지됩니다.

PyTorch의 DistributedDataParallel 클래스의 join API를 사용하여 데이터 병렬 랭크마다 배치 수가 다른 경우를 처리하는 경우에도 동일한 TP_GROUP 랭크의 배치 수가 동일한지 확인해야 합니다. 그렇지 않으면 모듈의 분산 실행에 사용되는 통신 콜렉티브가 중단될 수 있습니다. join API를 사용하는 한, TP_GROUP이 서로 다른 랭크라도 배치 수가 다를 수 있습니다.

모델을 체크포인트로 사용하고 텐서 병렬화를 사용하려면 다음 사항을 고려하세요.

텐서 병렬성을 사용하는 모델을 체크포인팅하는 방법에 대한 자세한 내용은 단원을 참조하세요.분산 모델 체크포인트 지정