PyTorch - Amazon SageMaker
Modelos PyTorch com transformadores Hugging Face

PyTorch

Traga seu próprio modelo do PyTorch para o SageMaker e execute o trabalho de treinamento com o SageMaker Training Compiler.

Modelos PyTorch com transformadores Hugging Face

Os modelos PyTorch com Hugging Face Transformers são baseados na API torch.nn.Module do PyTorch. O Hugging Face Transformers também fornece classes de Trainer e modelos pré-treinados para PyTorch para ajudar a reduzir o esforço de configuração de modelos de processamento de linguagem natural (NLP). Depois de preparar seu script de treinamento, você pode iniciar um trabalho de treinamento usando o SageMaker PyTorch ou o estimador HuggingFace com a configuração do SageMaker Training Compiler quando passar para o próximo tópico em Habilitar o SageMaker Training Compiler.

dica

Quando criar um tokenizador para um modelo de PNL usando Transformers em seu script de treinamento, certifique-se de usar uma forma de tensor de entrada estática especificando padding='max_length'. Não use padding='longest' porque o preenchimento da sequência mais longa do lote pode alterar a forma do tensor de cada lote de treinamento. A forma de entrada dinâmica pode acionar a recompilação do modelo e aumentar o tempo total de treinamento. Para obter mais informações sobre as opções de preenchimento dos tokenizadores Transformers, consulte Preenchimento e truncamento na documentação de Hugging Face Transformers.

Modelos de linguagem grandes usando a classe Trainer de Hugging Face Transformers

Se você usa a classe Trainer da biblioteca transformers, não precisa fazer nenhuma alteração adicional em seu script de treinamento. O SageMaker Training Compiler compila automaticamente seu modelo Trainer se você ativá-lo por meio da classe estimador. O código a seguir mostra a forma básica de um script de treinamento do PyTorch com a API Hugging Face Trainer.

from transformers import Trainer, TrainingArguments

training_args=TrainingArguments(**kwargs)
trainer=Trainer(args=training_args, **kwargs)

Para treinamento em uma única GPU

Você não precisa alterar o código quando usar a classe transformers.Trainer.

Para treinamento distribuído

PyTorch v1.11.0 e versões posteriores

Para executar um treinamento distribuído com o SageMaker Training Compiler, adicione a função_mp_fn() a seguir em seu script de treinamento e agrupe a função main(). As chamadas de função _mp_fn(index) do runtime distribuído do SageMaker são redirecionadas para PyTorch (pytorchxla) para a função main() do seu script de treinamento. 

def _mp_fn(index):
    main()

Essa função aceita o argumento index para indicar a classificação da GPU atual no cluster para treinamento distribuído. Para encontrar mais exemplos de scripts, consulte os scripts de exemplo de modelagem da linguagem Hugging Face Transformers.

Para Transformers v4.17 e anteriores com PyTorch v1.10.2 e anteriores

O SageMaker Training Compiler usa um mecanismo alternativo para iniciar um trabalho de treinamento distribuído e você não precisa fazer nenhuma modificação em seu script de treinamento. Em vez disso, o SageMaker Training Compiler exige que você passe um script de inicialização de treinamento distribuído do SageMaker para o argumento entry_point e passe seu script de treinamento para o argumento hyperparameters no estimador Hugging Face do SageMaker.

Melhores práticas para usar o SageMaker Training Compiler com Trainer

  • Certifique-se de usar os otimizadores SyncFree definindo o argumento optim como adamw_torch_xla ao configurar Transformers.trainingArgument. Veja também Optimizer na documentação do Hugging Face Transformers.

  • Certifique-se de que a taxa de transferência do pipeline de processamento de dados seja maior do que o throughput do treinamento. Você pode ajustar os argumentos dataloader_num_workers e preprocessing_num_workers da classe transformers.TrainingArgument para fazer isso. Normalmente, eles precisam ser maiores ou iguais ao número de GPUs, mas menores que o número de CPUs.

Depois de concluir a adaptação do seu roteiro de treinamento, prossiga para Execute trabalhos de treinamento do PyTorch com o SageMaker Training Compiler.

Modelos de linguagem grandes usando PyTorch diretamente (sem a API Hugging Face Transformers Trainer)

Se você tiver um script de treinamento que usa o PyTorch diretamente, precisará fazer alterações adicionais no script de treinamento do PyTorch para implementar o PyTorch/XLA. Siga as instruções para modificar seu script para configurar adequadamente as primativas PyTorch/XLA.

Para treinamento em uma única GPU

  1. Importe as bibliotecas de otimização.

    import torch_xla
import torch_xla.core.xla_model as xm

  2. Altere o dispositivo de destino para XLA em vez de torch.device("cuda")

    device=xm.xla_device()

  3. Se você estiver usando a Precisão Mista Automática (AMP) do PyTorch, faça o seguinte:

    1. Substitua torch.cuda.amp pelo seguinte:

      import torch_xla.amp

    2. Substitua torch.optim.SGD e torch.optim.Adam por um dos seguintes:

      import torch_xla.amp.syncfree.Adam as adam
import torch_xla.amp.syncfree.SGD as SGD

    3. Substitua torch.cuda.amp.GradScaler pelo seguinte:

      import torch_xla.amp.GradScaler as grad_scaler

  4. Se você não estiver usando AMP, substitua optimizer.step() pelo seguinte:

    xm.optimizer_step(optimizer)

  5. Se estiver usando um carregador de dados distribuído, envolva seu carregador de dados na classe ParallelLoader do PyTorch/XLA:

    import torch_xla.distributed.parallel_loader as pl
parallel_loader=pl.ParallelLoader(dataloader, [device]).per_device_loader(device)

  6. Adicione mark_step no final do ciclo de treinamento quando não estiver usando parallel_loader:

    xm.mark_step()

  7. Para verificar seu treinamento, use o método de ponto de verificação do modelo PyTorch/XLA:

    xm.save(model.state_dict(), path_to_save)

Depois de concluir a adaptação do seu roteiro de treinamento, prossiga para Execute trabalhos de treinamento do PyTorch com o SageMaker Training Compiler.

Para treinamento distribuído

Além das alterações listadas na seção Para treinamento em uma única GPU anterior, adicione as seguintes alterações para distribuir adequadamente o workload entre as GPUs.

  1. Se estiver usando AMP, adicione all_reduce depois scaler.scale(loss).backward():

    gradients=xm._fetch_gradients(optimizer)
xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size())

  2. Se você precisar definir variáveis para local_ranks eworld_size, use um código semelhante ao seguinte:

    local_rank=xm.get_local_ordinal()
world_size=xm.xrt_world_size()

  3. Para qualquer world_size (num_gpus_per_node*num_nodes) maior que 1, defina uma amostra de treino que deve ser semelhante ao seguinte:

    import torch_xla.core.xla_model as xm

if xm.xrt_world_size() > 1:
    train_sampler=torch.utils.data.distributed.DistributedSampler(
        train_dataset,
        num_replicas=xm.xrt_world_size(),
        rank=xm.get_ordinal(),
        shuffle=True
    )

train_loader=torch.utils.data.DataLoader(
    train_dataset, 
    batch_size=args.batch_size,
    sampler=train_sampler,
    drop_last=args.drop_last,
    shuffle=False if train_sampler else True,
    num_workers=args.num_workers
)

  4. Faça as seguintes alterações para garantir que você use o parallel_loader fornecido pelo módulotorch_xla distributed . 

    import torch_xla.distributed.parallel_loader as pl
train_device_loader=pl.MpDeviceLoader(train_loader, device)

    As funções train_device_loader são como um carregador PyTorch normal da seguinte forma: 

    for step, (data, target) in enumerate(train_device_loader):
    optimizer.zero_grad()
    output=model(data)
    loss=torch.nn.NLLLoss(output, target)
    loss.backward()

    Com todas essas mudanças, você deve ser capaz de iniciar o treinamento distribuído com qualquer modelo do PyTorch sem a API do Transformer Trainer. Observe que essas instruções podem ser usadas tanto para várias GPUs de nó único quanto para várias GPUs de vários nós.

  5. Para PyTorch v1.11.0 e versões posteriores

    Para executar um treinamento distribuído com o SageMaker Training Compiler, adicione a função _mp_fn() a seguir em seu script de treinamento e agrupe a função main(). As chamadas de função _mp_fn(index) do runtime distribuído do SageMaker são redirecionadas para PyTorch (pytorchxla) para a função main() do seu script de treinamento. 

    def _mp_fn(index):
    main()

    Essa função aceita o argumento index para indicar a classificação da GPU atual no cluster para treinamento distribuído. Para encontrar mais exemplos de scripts, consulte os scripts de exemplo de modelagem da linguagem Hugging Face Transformers.

    Para Transformers v4.17 e anteriores com PyTorch v1.10.2 e anteriores

    O SageMaker Training Compiler usa um mecanismo alternativo para iniciar um trabalho de treinamento distribuído e exige que você passe um script inicializador de treinamento distribuído do SageMaker para o argumento entry_point e passe seu script de treinamento para o argumento hyperparameters no estimador Hugging Face do SageMaker.

Depois de concluir a adaptação do seu script de treinamento, prossiga para Execute trabalhos de treinamento do PyTorch com o SageMaker Training Compiler.

Melhores práticas para usar o SageMaker Training Compiler com PyTorch/XLA

Se quiser aproveitar o compilador de treinamento do SageMaker em seu script de treinamento nativo do PyTorch, talvez queira primeiro se familiarizar com o PyTorch em dispositivos XLA. As seções a seguir listam algumas das melhores práticas para habilitar o XLA para PyTorch.

nota

Esta seção de melhores práticas pressupõe que você use os seguintes módulos PyTorch/XLA:

import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl
Entenda o modo lazy em PyTorch/XLA

Uma diferença significativa entre o PyTorch/XLA e o PyTorch nativo é que o sistema PyTorch/XLA é executado no modo lazy, enquanto o PyTorch nativo é executado no modo eager. Os tensores no modo lazy são espaços reservados para construir o gráfico computacional até que sejam materializados após a conclusão da compilação e avaliação. O sistema PyTorch/XLA cria o gráfico computacional dinamicamente quando chamamos as APIs do PyTorch para criar a computação usando tensores e operadores. O gráfico computacional é compilado e executado quando xm.mark_step() é chamado explícita ou implicitamente porpl.MpDeviceLoader/pl.ParallelLoader, ou quando você solicita explicitamente o valor de um tensor, como, por exemplo, chamando loss.item() ou print(loss).

Minimize o número de compilações e execuções usando pl.MpDeviceLoader/pl.ParallelLoader e xm.step_closure

Para obter o melhor desempenho, lembre-se das formas possíveis de iniciar a compilação e as execuções, conforme descrito em Entenda o modo lazy em PyTorch/XLA e tente minimizar o número de compilações e execuções. Idealmente, apenas uma compilação e execução é necessária por iteração de treinamento e é iniciada automaticamente pelo pl.MpDeviceLoader/pl.ParallelLoader. O MpDeviceLoader é otimizado para XLA e sempre deve ser usado, se possível, para obter o melhor desempenho. Durante o treinamento, talvez você queira examinar alguns resultados intermediários, como valores de perda. Nesse caso, a impressão de tensores lazy deve ser empacotada usando-se xm.add_step_closure() para evitar compilações e execuções desnecessárias.

Use AMP e otimizadores syncfree

O treinamento no modo Automatic Mixed Precision (AMP) acelera significativamente sua velocidade de treinamento ao aproveitar os núcleos tensores das GPUs NVIDIA. O SageMaker Training Compiler fornece otimizadores syncfree que são otimizados para XLA para melhorar o desempenho do AMP. Atualmente, os três otimizadores syncfree a seguir estão disponíveis e devem ser usados, se possível, para obtenção do melhor desempenho.

torch_xla.amp.syncfree.SGD
torch_xla.amp.syncfree.Adam
torch_xla.amp.syncfree.AdamW

Esses otimizadores syncfree devem ser combinados para escalonamento/desescalonamento de gradiente torch_xla.amp.GradScaler.

dica

A partir do PyTorch 1.13.1, o SageMaker Training Compiler melhora o desempenho ao permitir que o PyTorch/XLA substitua automaticamente os otimizadores (como SGD, Adam, AdamW) em torch.optim ou transformers.optimization com suas versões sem sincronização em torch_xla.amp.syncfree (como torch_xla.amp.syncfree.SGD, torch_xla.amp.syncfree.Adam, torch_xla.amp.syncfree.AdamW). Você não precisa alterar as linhas de código nas quais define otimizadores em seu script de treinamento.