PyTorch 교육 스크립트 수정

이 섹션에서는 PyTorch 교육 스크립트를 수정하여 자동 파티셔닝과 수동 파티셔닝을 위한 SageMaker 모델 병렬화 라이브러리를 구성하는 방법을 알아봅니다.

참고

라이브러리에서 지원하는 PyTorch 버전을 찾으려면 을 참조하십시오. 지원되는 프레임워크 및 AWS 리전

작은 정보

SageMaker 모델 병렬 처리 라이브러리와 함께 PyTorch 교육 스크립트를 사용하는 방법을 보여주는 end-to-end 노트북 예제는 을 참조하십시오. Amazon SageMaker 모델 병렬화 라이브러리 v1 예제

자동 파티셔닝은 기본적으로 활성화되어 있음을 유의하세요. 달리 지정하지 않는 한, 다음 스크립트는 자동 파티셔닝을 사용합니다.

다음과 같은 자동 분할 PyTorch

SageMaker의 모델 병렬 처리 라이브러리로 교육 스크립트를 실행하려면 다음과 같은 PyTorch 교육 스크립트 변경이 필요합니다.

1. smdistributed.modelparallel.torch.init()을 사용하여 라이브러리를 가져오고 초기화합니다.

2. smdistributed.modelparallel.torch.DistributedModel로 모델을 래핑합니다. 기본 nn.Module 객체의 forward 메서드에서 반환되는 모든 텐서는 모델 병렬 디바이스 간에 브로드캐스트되므로 통신 오버헤드가 발생하므로 호출 메서드 외부에서 필요하지 않은 텐서(예: 중간 활성화)는 반환되지 않아야 합니다.

   참고

   FP16 훈련의 경우 smdistributed.modelparallel.torch.model_create() 컨텍스트 관리자를 사용하여 모델을 래핑해야 합니다. 자세한 정보는 모델 병렬 처리를 사용한 FP16 훈련을 참조하세요.

3. smdistributed.modelparallel.torch.DistributedOptimizer를 이용해 옵티마이저를 래핑합니다.

   참고

   FP16 훈련의 경우 정적 또는 동적 손실 스케일링을 설정해야 합니다. 자세한 정보는 모델 병렬 처리를 사용한 FP16 훈련을 참조하세요.

4. 사용자 모델 대신 반환된 DistributedModel 객체를 사용하세요.

5. 순방향 및 역방향 로직을 Step Function에 넣고 이를 smdistributed.modelparallel.torch.step로 데코레이트하세요.

6. torch.cuda.set_device(smp.local_rank())를 통해 각 프로세스를 자체 디바이스로 제한합니다.

7. smp.step 호출 전에 .to() API를 사용하여 입력 텐서를 GPU로 이동합니다 (아래 예제 참조).

8. torch.Tensor.backward와 torch.autograd.backward를 DistributedModel.backward로 바꾸세요.

9. reduce_mean과 같은 StepOutput 메서드를 사용하여 마이크로배치의 출력값에 대해 후처리를 수행합니다.

10. 평가 단계가 있는 경우에도 마찬가지로 smp.step로 데코레이팅된 함수 안에 순방향 로직을 배치하고 StepOutput API를 사용하여 출력값을 후처리하세요.

11. DataLoader에서 drop_last=True로 설정합니다. 또는 배치 크기를 마이크로배치 수로 나눌 수 없는 경우 훈련 루프에서 배치를 수동으로 건너뛰어도 됩니다.

SageMaker의 모델 병렬화 라이브러리 API에 대한 자세한 내용은 API 설명서를 참조하십시오.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchnet.dataset import SplitDataset
from torchvision import datasets

import smdistributed.modelparallel.torch as smp

class GroupedNet(nn.Module):
    def __init__(self):
        super(GroupedNet, self).__init__()
        # define layers

    def forward(self, x):
        # define forward pass and return model outputs


# smdistributed: Define smp.step. Return any tensors needed outside.
@smp.step
def train_step(model, data, target):
    output = model(data)
    loss = F.nll_loss(output, target, reduction="mean")
    model.backward(loss)
    return output, loss


def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # smdistributed: Move input tensors to the GPU ID used by the current process,
        # based on the set_device call.
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        # Return value, loss_mb is a StepOutput object
        _, loss_mb = train_step(model, data, target)

        # smdistributed: Average the loss across microbatches.
        loss = loss_mb.reduce_mean()

        optimizer.step()

# smdistributed: initialize the backend
smp.init()

# smdistributed: Set the device to the GPU ID used by the current process.
# Input tensors should be transferred to this device.
torch.cuda.set_device(smp.local_rank())
device = torch.device("cuda")

# smdistributed: Download only on a single process per instance.
# When this is not present, the file is corrupted by multiple processes trying
# to download and extract at the same time
dataset = datasets.MNIST("../data", train=True, download=False)

# smdistributed: Shard the dataset based on data-parallel ranks
if smp.dp_size() > 1:
    partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())}
    dataset = SplitDataset(dataset, partitions=partitions_dict)
    dataset.select(f"{smp.dp_rank()}")

# smdistributed: Set drop_last=True to ensure that batch size is always divisible
# by the number of microbatches
train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)

model = GroupedNet()
optimizer = optim.Adadelta(model.parameters(), lr=4.0)

# smdistributed: Use the DistributedModel container to provide the model
# to be partitioned across different ranks. For the rest of the script,
# the returned DistributedModel object should be used in place of
# the model provided for DistributedModel class instantiation.
model = smp.DistributedModel(model)
optimizer = smp.DistributedOptimizer(optimizer)

train(model, device, train_loader, optimizer)

수동 분할은 다음과 같습니다. PyTorch

smp.partition 컨텍스트 관리자를 사용하여 특정 디바이스에 모듈을 배치합니다. smp.partition 컨텍스트에 배치되지 않은 모든 모듈은 default_partition에 배치됩니다. auto_partition가 False로 설정된 경우 default_partition을 제공해야 합니다. 특정 smp.partition 컨텍스트 내에서 생성된 모듈은 해당 파티션에 배치됩니다.

SageMaker의 모델 병렬화 라이브러리 API에 대한 자세한 내용은 API 설명서를 참조하십시오.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchnet.dataset import SplitDataset
from torchvision import datasets

import smdistributed.modelparallel.torch as smp

class GroupedNet(nn.Module):
    def __init__(self):
        super(GroupedNet, self).__init__()
        with smp.partition(0):
            # define child modules on device 0
        with smp.partition(1):
            # define child modules on device 1

    def forward(self, x):
        # define forward pass and return model outputs


# smdistributed: Define smp.step. Return any tensors needed outside.
@smp.step
def train_step(model, data, target):
    output = model(data)
    loss = F.nll_loss(output, target, reduction="mean")
    model.backward(loss)
    return output, loss


def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # smdistributed: Move input tensors to the GPU ID used by the current process,
        # based on the set_device call.
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        # Return value, loss_mb is a StepOutput object
        _, loss_mb = train_step(model, data, target)

        # smdistributed: Average the loss across microbatches.
        loss = loss_mb.reduce_mean()

        optimizer.step()

# smdistributed: initialize the backend
smp.init()

# smdistributed: Set the device to the GPU ID used by the current process.
# Input tensors should be transferred to this device.
torch.cuda.set_device(smp.local_rank())
device = torch.device("cuda")

# smdistributed: Download only on a single process per instance.
# When this is not present, the file is corrupted by multiple processes trying
# to download and extract at the same time
dataset = datasets.MNIST("../data", train=True, download=False)

# smdistributed: Shard the dataset based on data-parallel ranks
if smp.dp_size() > 1:
    partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())}
    dataset = SplitDataset(dataset, partitions=partitions_dict)
    dataset.select(f"{smp.dp_rank()}")

# smdistributed: Set drop_last=True to ensure that batch size is always divisible
# by the number of microbatches
train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)

model = GroupedNet()
optimizer = optim.Adadelta(model.parameters(), lr=4.0)

# smdistributed: Use the DistributedModel container to provide the model
# to be partitioned across different ranks. For the rest of the script,
# the returned DistributedModel object should be used in place of
# the model provided for DistributedModel class instantiation.
model = smp.DistributedModel(model)
optimizer = smp.DistributedOptimizer(optimizer)

train(model, device, train_loader, optimizer)

고려 사항

SageMaker의 모델 병렬화 라이브러리를 사용하여 PyTorch 교육 스크립트를 구성할 때는 다음 사항을 숙지해야 합니다.

• 글로벌 그래디언트 규범(예: LAMB 옵티마이저의 일부 변형 또는 글로벌 그래디언트 클리핑 등 전체 모델의 그라데이션 규범)에 의존하는 최적화 기법을 사용하는 경우 정확성을 위해 모델 파티션의 모든 표준을 수집해야 합니다. 라이브러리의 통신 기본 데이터 유형을 사용하여 이 작업을 수행할 수 있습니다.

• 모델에 있는 nn.Modules의 전달 메서드에 대한 모든 torch.Tensor 인수는 모듈 출력 계산에 사용해야 합니다. 즉, 라이브러리는 모듈 출력이 종속되지 않는 모듈에 대한 torch.Tensor 인수가 있는 경우를 지원하지 않습니다.

• smp.DistributedModel.backward() 호출에 대한 인수는 모든 모델 출력에 종속되어야 합니다. 즉, smp.DistributedModel.forward 호출에 공급되는 텐서의 계산에 사용되지 않는 smp.DistributedModel.backward 호출의 출력은 있을 수 없습니다.

• 코드에 torch.cuda.synchronize() 호출이 있는 경우 동기화 호출 직전에 torch.cuda.set_device(smp.local_rank())를 호출해야 할 수도 있습니다. 그렇지 않으면 디바이스 0에 불필요한 CUDA 컨텍스트가 생성되어 메모리를 불필요하게 소비하게 될 수 있습니다.

• 라이브러리는 서로 다른 디바이스에 nn.Modules을 배치하므로 모델 내 모듈이 smp.step 내부에서 수정된 글로벌 상태에 종속되지 않아야 합니다. 훈련 기간 동안 고정된 상태로 유지되거나 모든 프로세스에서 볼 수 있는 방식으로 smp.step 외부에서 수정된 모든 상태는 허용됩니다.

• 라이브러리를 사용할 때는 모델을 GPU로 옮길 필요가 없습니다 (예: model.to(device) 사용). 모델이 파티셔닝되기 전(첫 번째 smp.step 호출 전)에 모델을 GPU로 이동하려고 하면 이동 호출이 무시됩니다. 라이브러리는 순위에 할당된 모델 부분을 자동으로 해당 GPU로 이동합니다. 라이브러리를 이용하는 훈련이 시작되면 모델을 CPU로 옮겨 사용하지 마세요. 프로세스에서 보유한 파티션에 할당되지 않은 모듈의 경우 올바른 파라미터를 가질 수 없기 때문입니다. 모델 병렬 처리 라이브러리를 사용하여 모델을 학습시킨 후 라이브러리 없이 모델을 다시 학습하거나 추론에 사용하려는 경우, 체크포인트 API를 사용하여 전체 모델을 저장하고 일반 모듈로 다시 로드하는 것이 좋습니다. PyTorch

• 한 피드를 다른 피드로 출력하는 것과 같은 모듈 목록이 있는 경우 해당 목록을 nn.Sequential로 바꾸면 성능이 크게 향상될 수 있습니다.

• 가중치 업데이트(optimizer.step())는 전체 역방향 패스가 완료되고 그래디언트가 준비되는 시점이기 때문에 smp.step의 외부에서 실행해야 합니다. 모델 및 데이터 병렬성을 갖춘 하이브리드 모델을 사용하는 경우 이 시점에서 기울기도 모두 완료될 수 있습니다. AllReduce

• 라이브러리를 데이터 병렬 처리와 함께 사용하는 경우 단계에 참여하지 AllReduce 않는 순위를 기다리는 시간이 발생하지 않도록 모든 데이터 병렬 랭크의 배치 수가 동일한지 확인하십시오.

• ml.p4d 인스턴스 유형(예: ml.p4d.24xlarge)을 사용하여 훈련 작업을 시작하는 경우 데이터 로더 변수 num_workers=0로 설정해야 합니다. 예를 들어, DataLoader를 다음과 같이 정의할 수 있습니다.

  dataloader = torch.utils.data.DataLoader(
              data,
              batch_size=batch_size,
              num_workers=0,
              pin_memory=True,
              drop_last=True,
              shuffle=shuffle,
          )

• smp.step에 대한 입력은 DataLoader에서 생성된 모델 입력이어야 합니다. 이는 smp.step가 배치 차원을 따라 입력 텐서를 내부적으로 분할하고 파이프라인을 생성하기 때문입니다. 즉, DataLoader를 smp.step 함수에 전달하여 내부에서 모델 입력을 생성하는 것은 작동하지 않습니다.

  예를 들어, DataLoader를 다음과 같이 정의하면

  train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)

  train_loader에서 생성된 모델 입력에 액세스하여 smp.step로 데코레이팅된 함수에 전달해야 합니다. smp.step함수에 train_loader를 직접 전달하지 마세요.

  def train(model, device, train_loader, optimizer):
      model.train()
      for batch_idx, (data, target) in enumerate(train_loader):
          ...
          _, loss_mb = train_step(model, data, target)
          ...
  
  @smp.step
  def train_step(model, data, target):
      ...
      return output, loss

• smp.step에의 입력 텐서는 .to() API를 사용하여 현재 디바이스로 이동해야 하며, 이는 torch.cuda.set_device(local_rank()) 호출 후에 이루어져야 합니다.

  예를 들어 다음과 같이 train 함수를 정의할 수 있습니다. 이 함수는 입력 텐서를 사용하여 train_step을 호출하기 전에 .to() API를 사용하여 현재 디바이스에 data와 target을 추가합니다.

  def train(model, device, train_loader, optimizer):
      model.train()
      for batch_idx, (data, target) in enumerate(train_loader):
          # smdistributed: Move input tensors to the GPU ID used by the current process,
          # based on the set_device call.
          data, target = data.to(device), target.to(device)
          optimizer.zero_grad()
          # Return value, loss_mb is a StepOutput object
          _, loss_mb = train_step(model, data, target)
  
          # smdistributed: Average the loss across microbatches.
          loss = loss_mb.reduce_mean()
  
          optimizer.step()

  위 train 함수에서 이 smp.set로 데코레이팅된 함수의 입력 텐서는 현재 디바이스로 이동되었습니다. 모델을 현재 디바이스로 옮길 필요는 없습니다. 라이브러리는 순위에 할당된 모델 부분을 자동으로 해당 GPU로 이동합니다.

  @smp.step
  def train_step(model, data, target):
      output = model(data)
      loss = F.nll_loss(output, target, reduction="mean")
      model.backward(loss)
      return output, loss

지원되지 않는 프레임워크 기능

의 모델 병렬화 PyTorch 라이브러리에서는 다음 기능을 지원하지 않습니다. SageMaker

• 네이티브 PyTorch DDP와 함께 데이터 병렬화를 사용하는 경우 라이브러리는 torch.nn.parallel.DistributedDataParallel래퍼 모듈을 지원하지 않습니다. 라이브러리는 파라미터 브로드캐스트 및 그래디언트를 포함하여 PyTorch DDP와의 통합을 내부적으로 관리합니다. AllReduce 라이브러리를 사용할 때 모듈 버퍼는 훈련 시작 시 한 번만 브로드캐스트됩니다. 모델에 각 단계에서 데이터 병렬 그룹 간에 동기화해야 하는 모듈 버퍼가 있는 경우 smp.get_dp_process_group()를 통해 얻을 수 있는 프로세스 그룹을 이용하는 torch.distributed API를 통해 그렇게 할 수 있습니다.

• 혼합 정밀도 훈련의 경우 apex.amp 모듈은 지원되지 않습니다. 라이브러리를 자동 혼합 정밀도와 함께 사용하는 권장 방법은 torch.cuda.amp를 이용하는 것입니다. 단, torch에서의 구현 대신 smp.amp.GradScaler를 사용하는 경우는 예외입니다.

• torch.jit.ScriptModules와 ScriptFunctions는 smp.DistributedModel에서 지원되지 않습니다.

• apex : apex의 FusedLayerNorm, FusedAdam, FusedLAMB, FusedNovoGrad는 지원되지 않습니다. 대신 smp.optimizers 및 smp.nn API를 통해 라이브러리로 이들을 구현할 수 있습니다.