Fragmentação de estado do otimizador - Amazon SageMaker

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á.

Fragmentação de estado do otimizador

A fragmentação de estado do otimizador é uma técnica útil de economia de memória que fragmenta o estado do otimizador (o conjunto de pesos que descreve o estado do otimizador) em grupos de dispositivos paralelos de dados. Você pode usar a fragmentação de estado do otimizador sempre que usar um otimizador com estado (como Adam) ou um FP16 otimizador (que armazena ambos FP16 e FP32 cópias dos parâmetros).

nota

A fragmentação de estado do otimizador está disponível PyTorch na biblioteca de paralelismo de SageMaker modelos v1.6.0 e versões posteriores.

Como usar a fragmentação de estado do otimizador

Você pode ativar a fragmentação de estado do otimizador definindo "shard_optimizer_state": True na configuração modelparallel.

Quando esse recurso é ativado, a biblioteca particiona o conjunto de parâmetros do modelo com base no grau de paralelismo de dados. Os gradientes correspondentes à i-ésima partição são reduzidos somente na i-enésima classificação paralela de dados. No final da primeira chamada para uma função decoradora smp.step, o otimizador encapsulado por smp.DistributedOptimizer redefine seus parâmetros para serem limitados apenas aos parâmetros correspondentes à partição da classificação paralela de dados atual. Os parâmetros redefinidos são chamados de parâmetros virtuais e compartilham o armazenamento subjacente com os parâmetros originais. Durante a primeira chamada para optimizer.step, os estados do otimizador são criados com base nesses parâmetros redefinidos, que são fragmentados por causa da partição original. Após a atualização do otimizador, a AllGather operação (como parte da optimizer.step chamada) é executada nas classificações paralelas de dados para obter estados de parâmetros consistentes.

dica

A fragmentação de estado do otimizador pode ser útil quando o grau de paralelismo de dados é maior que 1 e o modelo tem mais de um bilhão de parâmetros.

O grau de paralelismo de dados é calculado por (processes_per_host * instance_count / pipeline_parallel_degree), e a função smp.dp_size() lida com o dimensionamento em segundo plano.

Configurar um SageMaker PyTorch estimador

mpi_options = { "enabled" : True, "processes_per_host" : 8, # 8 processes "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none " } smp_options = { "enabled":True, "parameters": { "microbatches": 4, "pipeline_parallel_degree": 2, # alias for "partitions" "placement_strategy": "cluster", "tensor_parallel_degree": 2, # tp over 2 devices "ddp": True, "shard_optimizer_state": True } }

Adapte seu roteiro PyTorch de treinamento

Consulte Adaptar seu script PyTorch de treinamento na seção Paralelismo do tensor combinado com paralelismo do pipeline. Não há nenhuma modificação adicional necessária para o script.