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 otimizador FP16 (que armazena cópias dos parâmetros em FP16 e FP32).

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.