Divisão automatizada com PyTorch Divisão manual com PyTorch Considerações Recursos de framework incompatíveis

Modificar um script PyTorch de treinamento

Nesta seção, você aprende a modificar scripts de PyTorch treinamento para configurar a biblioteca de paralelismo de SageMaker modelos para particionamento automático e particionamento manual.

nota

Para descobrir quais PyTorch versões são suportadas pela biblioteca, consulteFrameworks compatíveis e Regiões da AWS.

dica

Para exemplos de end-to-end cadernos que demonstram como usar um script de PyTorch treinamento com a biblioteca de paralelismo de SageMaker modelos, consulte. Exemplos da biblioteca de paralelismo de SageMaker modelos da Amazon v1

Observe que o particionamento automático está habilitado por padrão. A menos que seja especificado de outra forma, os scripts a seguir utilizam autoparticionamento.

Divisão automatizada com PyTorch

As seguintes alterações no script de treinamento são necessárias para executar um script de PyTorch treinamento com a biblioteca SageMaker de paralelismo de modelos da:

Importe e inicialize a biblioteca com o smdistributed.modelparallel.torch.init().
Empacote o modelo com smdistributed.modelparallel.torch.DistributedModel. Esteja ciente de que quaisquer tensores retornados pelo método forward do objeto nn.Module subjacente serão transmitidos para os dispositivos de paralelismo de modelo, incorrendo em sobrecarga de comunicação. Portanto, quaisquer tensores que não são necessários fora do método de chamada (como ativações intermediárias) não devem ser retornados.

nota
Para o treinamento do FP16, você precisa usar o gerenciador de contexto smdistributed.modelparallel.torch.model_creation() para empacotar o modelo. Para ter mais informações, consulte Treinamento FP16 com paralelismo de modelos.
Empacotar o otimizador com smdistributed.modelparallel.torch.DistributedOptimizer.

nota
Para o treinamento do FP16, você precisa configurar a escalabilidade de perda estática ou dinâmica. Para ter mais informações, consulte Treinamento FP16 com paralelismo de modelos.
Use o objeto DistributedModel retornado em vez de um modelo de usuário.
Coloque a lógica para frente e para trás em uma step function e decore-a com smdistributed.modelparallel.torch.step.
Restrinja cada processo ao seu próprio dispositivo por meio de torch.cuda.set_device(smp.local_rank()).
Mova os tensores de entrada para a GPU usando a API do .to() antes da chamada do smp.step (veja o exemplo abaixo).
Substitua o torch.Tensor.backward e o torch.autograd.backward pelo DistributedModel.backward.
Execute o pós-processamento nas saídas em microlotes usando métodos StepOutput como reduce_mean.
Se houver uma etapa de avaliação, coloque logicamente a frente (forward) dentro de uma função decorada com smp.step e processe os resultados usando a API do StepOutput.
Defina drop_last=True em DataLoader. Como alternativa, pule manualmente um lote no ciclo de treinamento se o tamanho do lote não for divisível pelo número de microlotes.

Para saber mais sobre a API SageMaker da biblioteca de paralelismo de modelos, consulte a documentação da 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)

Divisão manual com PyTorch

Use gerenciadores de contexto do smp.partition para colocar módulos em dispositivos específicos. Qualquer operação não colocada em qualquer contexto smp.partition é colocada no default_partition. O default_partition precisa ser fornecido se o auto_partition estiver definido como False. Os módulos criados em um contexto smp.partition específico são colocados na partição correspondente.

Para saber mais sobre a API SageMaker da biblioteca de paralelismo de modelos, consulte a documentação da 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)

Considerações

Ao configurar um script de PyTorch treinamento usando a biblioteca SageMaker de paralelismo de modelos da, você deve estar ciente do seguinte:

Se você estiver usando uma técnica de otimização que depende de normas de gradiente globais, por exemplo, a norma de gradiente de todo o modelo, como algumas variantes do otimizador LAMB ou o clipping global de gradiente, é necessário reunir todas as normas nas partições do modelo para garantir a correção. Você pode usar os tipos de dados básicos de comunicação da biblioteca para fazer isso.
Todos os argumentos do torch.Tensor para os métodos diretos do nn.Modules em seu modelo devem ser usados no cálculo da saída do módulo. Em outras palavras, a biblioteca não suporta o caso em que há um argumento do torch.Tensor para um módulo do qual a saída do módulo não depende.
O argumento para a chamada smp.DistributedModel.backward() deve depender de todas as saídas do modelo. Em outras palavras, não pode haver uma saída da chamada smp.DistributedModel.forward que não seja usada no cálculo do tensor que é alimentado na chamada smp.DistributedModel.backward.
Se houver chamadas torch.cuda.synchronize() em seu código, talvez seja necessário ligar torch.cuda.set_device(smp.local_rank()) imediatamente antes da chamada de sincronização. Caso contrário, contextos CUDA desnecessários podem ser criados no dispositivo 0, o que consumirá desnecessariamente memória.
Dado que a biblioteca coloca nn.Modules em dispositivos diferentes, os módulos no modelo não devem depender de nenhum estado global que seja modificado dentro de smp.step. Qualquer estado que permaneça constante ao longo do treinamento, ou que seja modificado fora de smp.step de uma maneira que seja visível para todos os processos, é permitido.
Você não precisa mover o modelo para a GPU (por exemplo, usandomodel.to(device)) ao usar a biblioteca. Se você tentar mover o modelo para a GPU antes que o modelo seja particionado (antes da primeira chamada smp.step), a chamada de movimentação será ignorada. A biblioteca move automaticamente a parte do modelo atribuída a uma classificação para sua GPU. Quando o treinamento com a biblioteca começar, não mova o modelo para a CPU e o utilize, pois ele não terá os parâmetros corretos para módulos não atribuídos à partição mantida pelo processo. Se você quiser treinar novamente um modelo ou usá-lo para inferência sem a biblioteca depois de treiná-lo usando a biblioteca de paralelismo de modelos, a maneira recomendada é salvar o modelo completo usando nossa API de ponto de verificação e carregá-lo de volta em um módulo normal. PyTorch
Se você tem uma lista de módulos de forma que a saída de um alimenta o próximo, substituir essa lista por nn.Sequential pode melhorar significativamente a performance.
A atualização dos pesos (optimizer.step()) precisa ocorrer fora de smp.step, porque é nesse momento que toda a passagem de retropropagação é concluída e os gradientes estão prontos. Ao usar um modelo híbrido com paralelismo de modelos e dados, neste ponto, também é garantido o AllReduce término dos gradientes.
Ao usar a biblioteca em combinação com o paralelismo de dados, certifique-se de que o número de lotes em todas as classificações paralelas de dados seja o mesmo para que você AllReduce não fique esperando por uma classificação que não esteja participando da etapa.
Se você iniciar um trabalho de treinamento usando um tipo de instância ml.p4d (como ml.p4d.24xlarge), é necessário definir a variável do carregador de dados num_workers=0. Por exemplo, é possível definir o seu DataLoader da seguinte forma.
```
dataloader = torch.utils.data.DataLoader(
            data,
            batch_size=batch_size,
            num_workers=0,
            pin_memory=True,
            drop_last=True,
            shuffle=shuffle,
        )
```
As entradas para smp.step devem ser as entradas do modelo geradas pelo DataLoader. Isso ocorre porque divide smp.step internamente os tensores de entrada ao longo da dimensão do lote e os canaliza. Isso significa que passar a DataLoader si mesmo para a função smp.step para gerar as entradas internas do modelo não funciona.

Por exemplo, se definir um DataLoader da seguinte forma.
```
train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)
```
Você deve acessar as entradas do modelo geradas train_loader e passá-las para uma função smp.step decorada. Não passe train_loader diretamente para a função smp.step.
```
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
```

Os tensores de entrada smp.step devem ser movidos para o dispositivo atual usando a API do .to(), que deve ocorrer após a chamada torch.cuda.set_device(local_rank()).

Por exemplo, é possível definir a função train da seguinte forma. Essa função adiciona data e target ao dispositivo atual usando a API do .to() antes de usar esses tensores de entrada para chamar train_step.


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()

Os tensores de entrada para essa função smp.set decorada foram movidos para o dispositivo atual na função train acima. O modelo não precisa ser movido para o dispositivo atual. A biblioteca move automaticamente a parte do modelo atribuída a uma classificação para sua 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

Recursos de framework incompatíveis

Os seguintes PyTorch recursos não são compatíveis com a biblioteca de SageMaker paralelismo de modelos da:

Se você usa paralelismo de dados com o PyTorch DDP nativo, o módulo torch.nn.parallel.DistributedDataParallelwrapper não é suportado pela biblioteca. A biblioteca gerencia internamente a integração com o PyTorch DDP, incluindo transmissão de parâmetros e gradiente. AllReduce Ao utilizar a biblioteca, os buffers do módulo são transmitidos apenas uma vez no início do treinamento. Se seu modelo tiver buffers de módulo que precisam ser sincronizados entre grupos paralelos de dados em cada etapa, você pode fazer isso por meio da API do torch.distributed, usando o grupo de processos que pode ser obtido por meio de smp.get_dp_process_group().
Para treinamento de precisão mista, o módulo apex.amp não tem suporte. A maneira recomendada de usar a biblioteca com precisão mista automática é usar torch.cuda.amp, com exceção do uso de smp.amp.GradScaler em vez da implementação em torch.
torch.jit.ScriptModules e ScriptFunctions não têm suporte de smp.DistributedModel.
apex : FusedLayerNorm, FusedAdam, FusedLAMB e FusedNovoGrad do apex não têm suporte. Em vez disso, você pode usar as implementações de biblioteca por meio de smp.optimizers e APIs do smp.nn.

Atenção O Javascript está desativado ou não está disponível no seu navegador.

Para usar a documentação da AWS, o Javascript deve estar ativado. Consulte as páginas de Ajuda do navegador para obter instruções.

Convenções do documento

TensorFlow

Etapa 2: iniciar um trabalho de treinamento