Treine modelos gigantescos com dimensionamento quase linear usando paralelismo de dados fragmentados no Amazon SageMaker

Na busca por precisão superior, os modelos de aprendizado profundo em áreas como processamento de linguagem natural e visão computacional cresceram significativamente em tamanho nos últimos anos, frequentemente contados em dezenas a centenas de bilhões de parâmetros. Treinar esses modelos gigantescos é desafiador e requer estratégias de distribuição complexas. Cientistas de dados e engenheiros de aprendizado de máquina estão constantemente procurando a melhor maneira de otimizar sua computação de treinamento, mas estão lutando com a sobrecarga de comunicação que pode aumentar junto com o tamanho geral do cluster.

É por isso que lançamos recentemente sparalelismo de dados harded on Amazon Sage Maker, uma nova técnica de treinamento distribuído de economia de memória no Biblioteca de modelos paralelos (SMP) do SageMaker. O paralelismo de dados compartilhados foi desenvolvido especificamente para modelos de escala extrema e usa a Amazon internamente MiCS tecnologia sob o capô, um esforço científico para minimizar a escala de comunicação, reduzindo a dispendiosa sobrecarga de comunicação enraizada na coleta de parâmetros e na sincronização de gradiente. Com um modelo GPT-30 de parâmetro 2B com comprimento de sequência 2048, esse novo recurso alcançou 141 TFLOPs, uma velocidade de 39.7% em comparação com o DeepSpeed ​​ZeRO-3. Para um modelo 10B GPT-2 com comprimento de sequência 512, esse novo recurso também alcançou 564 amostras por segundo, uma velocidade de 13.9% em comparação com o Fully Sharded Data Parallel (FSDP) do PyTorch. Lembre-se de que no treinamento de modelos gigantescos, cada porcentagem de aceleração se traduz em economia de dinheiro e ganho de produtividade em sua equipe.

Nesta postagem do blog, primeiro examinaremos mais de perto os principais diferenciais do paralelismo de dados fragmentados e quando usá-lo. Em seguida, você aprenderá a treinar um modelo GPT-30 de parâmetro 2B no SageMaker com facilidade com esse novo recurso. Por fim, compararemos o desempenho com outras opções de código aberto, superando notavelmente o DeepSpeed ​​ZeRO em até 39.7% em 256 GPUs.

Como o paralelismo de dados fragmentados funciona e quando usá-lo

Antes de introduzirmos o paralelismo de dados fragmentados, vejamos sua família de técnicas mais ampla. Abordagens de treinamento distribuído recentes para grandes modelos mudaram para um paradigma em que os parâmetros do modelo, gradientes e estados do otimizador são compartilhados entre nós de dados paralelos. Ao contrário do paralelismo de pipeline, que tem a complexidade inata de escolher camadas para particionar entre dispositivos, especialmente quando sua estrutura não suporta divisão automatizada de modelos, esse paradigma preserva elegantemente a simplicidade do paralelismo de dados, enquanto remove a restrição do paralelismo de dados em que um modelo deve caber em uma única GPU.

Em estruturas existentes que se enquadram nesse paradigma, notadamente DeepSpeed ​​ZeRO-3 e FSDP do PyTorch upstream do FairScale, os estados do modelo são fragmentados em todos os GPUs, uma estratégia que reduz o consumo de memória em cada GPU ao custo de incorrer em grande sobrecarga de comunicação que aumenta com o tamanho do cluster e, portanto, faz com que a escalabilidade diminua significativamente em escala. Em contraste, o paralelismo de dados fragmentados nos estados do modelo de partições de biblioteca SMP em um ciente de escala maneira, particionando cada réplica de estados de modelo apenas dentro um subconjunto de GPUs.

Vamos olhar mais de perto o particionamento de modelo com reconhecimento de escala em MiCS, a tecnologia central por trás de dados fragmentados paralelos. A intuição por trás desse design é que o particionamento de estados de treinamento em todo o grupo de dados paralelos pode não ser necessário para treinar um modelo com dezenas de bilhões de parâmetros. Por exemplo, 8 GPUs V100 (32 GB cada) são suficientes para manter a réplica de estados do modelo de um modelo de parâmetro 10B que precisa de cerca de 200 GB de memória ao treinar com o otimizador Adam usando precisão mista. Ao limitar uma réplica completa dos estados do modelo no menor subconjunto de GPUs, podemos reduzir efetivamente a escala de sobrecarga de comunicação em comparação com DeepSpeed ​​e PyTorch FSDP. O paralelo de dados fragmentados também aproveita outras técnicas em MiCS, como comunicação hierárquica e sincronização de gradiente de 2 saltos. Para mais informações, confira Dimensionamento quase linear de treinamento de modelo gigantesco na AWS or MiCS: dimensionamento quase linear para treinamento de modelo gigantesco em nuvem pública.

Agora, como você sabe quando escolher dados fragmentados paralelos em vez de outras técnicas de treinamento distribuído? A regra geral é que, se o seu modelo tiver menos de 1 bilhão de parâmetros e puder caber na memória da GPU, Biblioteca paralela de dados SageMaker or Compilador de treinamento SageMaker pode ser suficiente para você. Se você possui modelos maiores de linguagem ou visão computacional, nossa sugestão é treiná-lo com a técnica de paralelismo de dados sharded combinada com checkpoint de ativação e descarga de ativação na biblioteca paralela do modelo SageMaker primeiro, antes de outras técnicas, como paralelismo tensorial ou paralelismo de pipeline.

Usar paralelismo de dados fragmentados para treinar GPT-2 no Amazon SageMaker

Vamos agora aprender como treinar um modelo GPT-2 com dados fragmentados em paralelo, com o SMP encapsulando a complexidade para você. este caderno tutorial completo orienta você por todo o processo, desde o processamento de dados, definição e envio de trabalhos de treinamento até o monitoramento de logs de treinamento. O que se segue é uma breve visão geral, destacando as principais etapas para usar esse recurso.

1. iniciar

O paralelismo de dados fragmentados está disponível no PyTorch v1.12.0+ e funciona com FP16 e BF16. A maneira mais fácil de usar a biblioteca SMP é por meio de um AWS Deep Learning Container for PyTorch pré-criado. No entanto, se você quiser trazer seu próprio contêiner Docker, consulte Crie seu próprio contêiner do Docker com a biblioteca paralela de modelos distribuídos do SageMaker. Para começar, siga Modificar um script de treinamento do PyTorch para adaptar as APIs dos SMPs em seu script de treinamento. Nesta seção, destacamos apenas algumas etapas principais com trechos de código do script de treinamento pronto para uso train_gpt_simple.py. Você pode seguir os comentários no script e Documento API para saber mais sobre onde as APIs SMP são usadas.

Primeiro, importe e inicialize a biblioteca chamando smdistributed.modelparallel.torch.init() no início do roteiro de treinamento:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

Em segundo lugar, envolva o modelo a ser particionado com smdistributed.modelparallel.torch.DistributedModel e use o retorno DistributedModel objeto daqui para frente:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel(model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation)

Envolva o otimizador com smdistributed.modelparallel.torch.DistributedOptimizer para salvar e carregar estados do otimizador.

from torch import optim

optimizer = optim.Adam(
    param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
)

optimizer = smp.DistributedOptimizer(
        optimizer, 
        static_loss_scale=None, 
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
        )

Coloque a lógica para frente e para trás em uma função de etapa e decore-a com smdistributed.modelparallel.torch.step.  Qualquer cálculo definido dentro do smp.step-decorated função é executada de forma distribuída.

@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    model.backward(loss)

    return loss

@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    
    return loss

2. Prepare o conjunto de dados

Usamos o openwebtext é o conjunto de dados que usamos neste exemplo. O notebook usa o script data_prep_512.py para baixar e pré-processar o conjunto de dados. Você também pode treinar com outros conjuntos de dados modificando data_pipeline.py. Ao lidar com grandes conjuntos de dados e modelos, você pode acelerar o trabalho de treinamento usando dados armazenados em Amazon FSx para Lustre, que fornece um sistema de arquivos de alto desempenho integrado nativamente com Serviço de armazenamento simples da Amazon (S3). Consulte as instruções de Configurar canal de entrada de dados para usar o Amazon FSx for Lustre para obter orientação sobre como configurar um sistema de arquivos FSx Lustre como canal de entrada de dados.

3. Inicie os trabalhos de treinamento

Esta etapa pressupõe que você já modificou seu script de treinamento e preparou o conjunto de dados conforme mencionado nas seções anteriores. Para habilitar o paralelismo de dados fragmentados, basta definir o sharded_data_parallel_degree no Estimador PyTorch. Neste tutorial, definimos sharded_data_parallel_degree=128 e instace_count=32 para nós p4d.24xlarge, o que indica que os estados do modelo serão fragmentados em 128 GPUs entre o total de 256 GPUs. Com base nesse valor selecionado, o SMP definirá automaticamente o grau de paralelo de dados para 2 (porque 256/128=2), o que significa que teremos duas réplicas para paralelismo de dados. Uma regra geral para escolher um valor ideal para sharded_data_parallel_degree é adicionar mais um nó ao grupo de compartilhamento por cada 3B de parâmetros do modelo. Neste tutorial, o tamanho do nosso modelo é 30B, portanto, devemos usar pelo menos 10 nós para fragmentação. E como 16 nós (128 GPUs) é a menor potência de 2 acima do limite, definimos sharded_data_parallel_degree=128.

Para checkpoints, também fornecemos um conjunto de utilitários de checkpoints em sharded_data_parallel_checkpoint.py , incluindo um utilitário para reconstruir o state_dict para casos de uso avançados. Finalmente, podemos iniciar um trabalho de treinamento distribuído chamando fit() no Estimator.

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    instance_type="ml.p4d.24xlarge",
    source_dir=os.getcwd(),
    volume_size=500,
    instance_count=32,
    distribution={
        "mpi": {
            "enabled": True,
            "processes_per_host": processes_per_host,
            "custom_mpi_options": mpioptions,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "ddp": True,
                    "skip_tracing": True,
                    "delayed_parameter_initialization": True,
                    "offload_activations": True,
                    "activation_loading_horizon": 4,
                    # To enable sharded data parallelism.
                    # Here we shard model states across 128 GPUs. 
                    "sharded_data_parallel_degree": 128, 
                    "fp16": False,
                    "bf16": True,
                    # This is to disable pipeline parallelism.
                    "partitions": 1,
                },
            }
        },
    },
    framework_version="1.12",
    py_version="py38",
    hyperparameters=hyperparameters,
    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,
    checkpoint_local_path=hyperparameters["checkpoint-dir"] if use_fsx else None,
    ...
)

smp_estimator.fit(inputs=data_channels)

4. Monitore os trabalhos de treinamento

Você pode acessar os logs de treinamento e acompanhar a utilização de GPU e memória em Amazon CloudWatch. Certifique-se de olhar para os logs de “algo-1” porque esse é o nó principal cujo fluxo de saída possui os logs de tarefas de treinamento de todas as instâncias.

Desempenho de benchmarking

Comparamos o paralelismo de dados fragmentados na biblioteca SMP em 16 e 32 nós p4d.24xlarge para o comprimento de sequência 512 e 2048, respectivamente. O modelo GPT30 de parâmetro 2B é configurado para usar uma largura oculta de 7168, 48 camadas e 64 cabeças. Você pode adotar exatamente a mesma configuração onde o comprimento da sequência é 2048 definindo model_config = "gpt2-30b" no caderno tutorial. Com essa configuração, o SMP alcançou 73.52 amostras por segundo, uma velocidade de 39.7% em comparação com o DeepSpeed ​​ZeRO-3. Se o tamanho do seu token for 500 bilhões, essa aceleração significa quase 367 horas de economia em nós p4d.24xlarge, o equivalente a mais de US$ 12,000 de orçamento economizado por treinamento! A tabela a seguir resume nossos resultados de benchmark.

1/ Para cada configuração de modelo, testamos diferentes recursos, estágios e configurações no DeepSpeed ​​ZeRO e escolhemos aquele que fornece o melhor rendimento como linha de base do DeepSpeed. O benchmark foi executado em Amazon Elastic Compute Nuvem (Amazônia EC2). 2/ Esses resultados contam com coletivos de comunicação aprimorados e otimizados para AWS que serão disponibilizados em breve. 3/ O tempo para treinar é projetado a partir da velocidade com base no número de tokens processados.

Em resumo, observamos uma taxa de transferência consistentemente mais alta com paralelismo de dados fragmentados no SMP quando comparado ao DeepSpeed ​​em uma variedade de modelos e configurações. Esse novo recurso também demonstrou uma melhor eficiência de memória em comparação com o DeepSpeed, permitindo que o SMP se ajuste a um tamanho de lote maior e reduza o nível de acúmulo de gradiente necessário para ajustar um tamanho de lote global específico.

Conclusão

Nesta postagem, apresentamos uma nova técnica de treinamento distribuído — paralelismo de dados fragmentados — e como ela acelera o treinamento gigantesco de modelos com dimensionamento quase linear no Amazon SageMaker. Também explicamos como treinar um modelo GPT-2 com a nova técnica seguindo este exemplo completo. Você pode seguir o Exemplos do Amazon SageMaker repositório GitHub para acompanhar todos os exemplos paralelos do modelo SageMaker ou participar do nosso próximo oficinas de treinamento distribuídas. Para saber mais sobre o paralelismo de dados fragmentados, consulte o documentação.

Sobre os autores

Emily Webber ingressou na AWS logo após o lançamento do SageMaker e vem tentando contar ao mundo sobre isso desde então! Além de criar novas experiências de ML para os clientes, Emily gosta de meditar e estudar o budismo tibetano.

Pode Karakus é Cientista Aplicado Sênior na AWS, otimizando o aprendizado profundo distribuído em larga escala na AWS. Seus interesses de pesquisa abrangem aprendizado profundo, otimização distribuída, sistemas distribuídos e teoria da informação. Fora do trabalho, ele gosta de andar de bicicleta, viajar, ler e aprender.

Rahul Huilgol é engenheiro de software sênior na AWS. Ele trabalha em sistemas distribuídos de aprendizado profundo, para tornar mais fácil e eficiente o treinamento de grandes modelos de aprendizado profundo na nuvem. Em seu tempo livre, ele gosta de fotografia, ciclismo e jardinagem.

Suhit Kodgule é um engenheiro de desenvolvimento de software do grupo de inteligência artificial da AWS que trabalha em estruturas de aprendizado profundo. Em seu tempo livre, ele gosta de caminhar, viajar e cozinhar.

Erin Ho é gerente de produto do AWS Deep Learning. Ela trabalha em produtos que facilitam para os clientes treinar modelos de aprendizado profundo na AWS. Para se divertir fora do trabalho, ela gosta de caminhar e esquiar.