Treinamento FP16 com paralelismo de modelos - Amazon SageMaker

Treinamento FP16 com paralelismo de modelos

Para o treinamento do FP16, aplique as seguintes modificações em seu script de treinamento e estimador:

nota

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

Adaptar seu script de treinamento do PyTorch

  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 wrapper do otimizador, consulte a API smdistributed.modelparallel.torch.DistributedOptimizer.

    O código a seguir é um exemplo de encapsulamento de um objeto otimizador Adadelta com escala dinâmica de perda para treinamento de 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 estimador SageMaker PyTorch

Adicione o parâmetro FP16 ("fp16") à configuração de distribuição para paralelismo do modelo ao criar um objeto estimador do SageMaker PyTorch. 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 treinamento do FP16 começa, o modelo e o otimizador são agrupados por FP16_Module e FP16_Optimizer respectivamente, que são versões smdistributed modificadas dos utilitários do Apex. FP16_Module converte o modelo em FP16 dtype e lida com a passagem direta no 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 um treinamento torch.optim.lr_scheduler e FP16, 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
)