모델 병렬 처리를 사용한 FP16 훈련 - Amazon SageMaker AI

기계 번역으로 제공되는 번역입니다. 제공된 번역과 원본 영어의 내용이 상충하는 경우에는 영어 버전이 우선합니다.

모델 병렬 처리를 사용한 FP16 훈련

FP16 훈련의 경우 훈련 스크립트와 예측기에 다음과 같은 수정 사항을 적용합니다.

참고

이 기능은 SageMaker 모델 병렬 처리 라이브러리 v1.10.0 이상에서 PyTorch에 사용할 수 있습니다.

PyTorch 훈련 스크립트 조정

  1. smdistributed.modelparallel.torch.model_creation() 컨텍스트 관리자를 사용하여 모델을 래핑합니다.

    # fp16_training_script.py

import torch
import smdistributed.modelparallel.torch as smp

with smp.model_creation(
    dtype=torch.float16 if args.fp16 else torch.get_default_dtype()
):
    model = ...
    작은 정보

    텐서 병렬 처리를 사용하는 경우 smp.model_creation 컨텍스트 관리자에 tensor_parallelism=smp.tp_size() > 1을 추가합니다. 이 줄을 추가하면 텐서 병렬 처리 활성화 여부를 자동으로 감지하는 데도 도움이 됩니다.

    with smp.model_creation(
    ... ,
    tensor_parallelism=smp.tp_size() > 1
):
    model = ...

  2. 옵티마이저를 smdistributed.modelparallel.torch.DistributedOptimizer로 래핑할 경우 static_loss_scaling 또는 dynamic_loss_scaling 인수를 설정합니다. 기본적으로 static_loss_scaling는 1.0으로 설정되고 dynamic_loss_scaling는 False로 설정됩니다. dynamic_loss_scale=True를 설정하면 dynamic_loss_args 인수를 통해 동적 손실 크기 조정 옵션을 사전으로 제공할 수 있습니다. 대부분의 경우 기본 옵션과 함께 동적 손실 크기 조정을 사용하는 것이 좋습니다. 옵티마이저 래퍼 함수에 대한 자세한 내용, 옵션 및 예제는 smdistributed.modelparallel.torch.DistributedOptimizer API를 참조하세요.

    다음 코드는 FP16 훈련을 위해 동적 손실 크기 조정으로 Adadelta 옵티마이저 객체를 래핑하는 예제입니다.

    optimizer = torch.optim.Adadelta(...)
optimizer = smp.DistributedOptimizer(
    optimizer,
    static_loss_scale=None,
    dynamic_loss_scale=True,
    dynamic_loss_args={
        "scale_window": 1000,
        "min_scale": 1,
        "delayed_shift": 2
    }
)

SageMaker PyTorch 예측기 구성

SageMaker PyTorch 예측기 객체를 생성할 때 모델 병렬 처리를 위한 분포 구성에 FP16 파라미터("fp16")를 추가합니다. 모델 병렬 처리용 구성 파라미터의 전체 목록은 smdistributed용 구성을 참조하세요.

from sagemaker.pytorch import PyTorch

smp_options = {
    "enabled": True,
    "parameters":  {
        "microbatches":  4,
        "pipeline_parallel_degree":  2,
        "tensor_parallel_degree":  2,
        ...,

        "fp16": True
    }
}

fp16_estimator = PyTorch(
    entry_point="fp16_training_script.py", # Specify your train script
    ...,

    distribution={
        "smdistributed": {"modelparallel": smp_options},
        "mpi": {...}
    }
)

fp16_estimator.fit(...)

FP16 훈련이 시작되면 모델 및 옵티마이저는 각각 FP16_Module 및 FP16_Optimizer가 래핑하며 이는 Apex utils의 수정된 smdistributed 버전입니다. FP16_Module은 모델을 FP16 dtype으로 변환하고 FP16에서 순방향 패스를 처리합니다.

작은 정보

optimizer.step 전에 clip_master_grads를 호출하여 그라디언트 클리핑을 적용합니다.

optimizer.clip_master_grads(max_norm)     # max_norm(float or int): max norm of the gradients
작은 정보

torch.optim.lr_scheduler 및 FP16 훈련을 사용할 때는 옵티마이저가 아닌 LR 스케줄러에 optimizer.optimizer를 전달해야 합니다. 다음 예제 코드를 참조하세요.

from torch.optim.lr_scheduler import StepLR

scheduler = StepLR(
    optimizer.optimizer if smp.state.cfg.fp16 else optimizer,
    step_size=1,
    gamma=args.gamma
)