FP16 Treinamento com paralelismo de modelos - SageMaker IA da Amazon

As traduções são geradas por tradução automática. Em caso de conflito entre o conteúdo da tradução e da versão original em inglês, a versão em inglês prevalecerá.

FP16 Treinamento com paralelismo de modelos

Para FP16 treinamento, aplique as seguintes modificações em seu roteiro de treinamento e estimador.

nota

Esse recurso está disponível PyTorch na biblioteca de paralelismo de SageMaker modelos v1.10.0 e versões posteriores.

Adapte seu roteiro PyTorch de treinamento

  1. Envolva seu modelo usando o gerenciador de contexto 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 = ...
    dica

    Se você estiver usando paralelismo de tensores, adicione tensor_parallelism=smp.tp_size() > 1 ao gerenciador de contexto smp.model_creation. Adicionar essa linha também ajuda a detectar automaticamente se o paralelismo do tensor está ativado ou não.

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

  2. Ao encapsular o otimizador com smdistributed.modelparallel.torch.DistributedOptimizer, defina o argumento static_loss_scaling ou dynamic_loss_scaling. Por padrão, static_loss_scaling está definido como 1.0 e dynamic_loss_scaling está definido como False. Se você definir dynamic_loss_scale=True, poderá alimentar as opções dinâmicas de escalabilidade de perda como um dicionário por meio do argumento dynamic_loss_args. Na maioria dos casos, recomendamos que você use a escala de perda dinâmica com as opções padrão. Para obter mais informações, opções e exemplos da função de encapsulamento do otimizador, consulte smdistributed.modelparallel.torch. DistributedOptimizer API.

    O código a seguir é um exemplo de encapsulamento de um objeto Adadelta otimizador com escala dinâmica de perda para treinamento. FP16 

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

Configurar um SageMaker PyTorch estimador

Adicione o FP16 parâmetro ("fp16") à configuração de distribuição para paralelismo do modelo ao criar um SageMaker PyTorch objeto estimador. Para obter uma lista completa dos parâmetros de configuração do paralelismo do modelo, consulte Parâmetros para 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(...)

Quando o FP16 treinamento começa, o modelo e o otimizador são agrupados por FP16_Module e, FP16_Optimizer respectivamente, que são smdistributed versões modificadas dos utilitários do Apex. FP16_Moduleconverte o modelo em FP16 dtype e lida com a passagem direta para dentro. FP16

dica

Você pode aplicar o recorte de gradiente chamando clip_master_grads antes de optimizer.step.

optimizer.clip_master_grads(max_norm)     # max_norm(float or int): max norm of the gradients
dica

Ao usar torch.optim.lr_scheduler e FP16 treinar, você precisa passar optimizer.optimizer para o agendador LR em vez do otimizador. Veja o exemplo de código a seguir.

from torch.optim.lr_scheduler import StepLR

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