A biblioteca paralela do modelo Amazon SageMaker agora acelera as cargas de trabalho PyTorch FSDP em até 20% | Amazon Web Services

A biblioteca paralela do modelo Amazon SageMaker agora acelera as cargas de trabalho PyTorch FSDP em até 20% | Amazon Web Services

Carimbo de data / hora: 22 de dezembro de 2023, 5h04
Nó Fonte: 3033172

O treinamento em modelos de linguagem grande (LLM) cresceu em popularidade no último ano com o lançamento de vários modelos populares, como Llama 2, Falcon e Mistral. Os clientes estão agora pré-treinando e ajustando LLMs que variam de 1 bilhão a mais de 175 bilhões de parâmetros para otimizar o desempenho do modelo para aplicações em todos os setores, desde saúde até finanças e marketing.

Treinar modelos de desempenho nesta escala pode ser um desafio. LLMs altamente precisos podem exigir terabytes de dados de treinamento e milhares ou até milhões de horas de tempo de computação do acelerador para atingir a precisão desejada. Para concluir o treinamento e lançar produtos em tempo hábil, os clientes contam com técnicas de paralelismo para distribuir essa enorme carga de trabalho por até milhares de dispositivos aceleradores. No entanto, estas técnicas de paralelismo podem ser difíceis de utilizar: diferentes técnicas e bibliotecas são compatíveis apenas com determinadas cargas de trabalho ou restritas a determinadas arquiteturas de modelo, o desempenho do treino pode ser altamente sensível a configurações obscuras e o estado da arte está a evoluir rapidamente. Como resultado, os profissionais de aprendizado de máquina devem passar semanas se preparando para escalar suas cargas de trabalho de LLM para grandes clusters de GPUs.

Neste post, destacamos novos recursos do Amazon Sage Maker biblioteca de modelo paralelo (SMP) que simplifica o processo de treinamento de modelos grandes e ajuda você a treinar LLMs com mais rapidez. Em particular, cobrimos a nova experiência de usuário simplificada da biblioteca SMP que se baseia em APIs PyTorch Fully Sharded Data Parallel (FSDP) de código aberto, funcionalidade paralela de tensor expandida que permite modelos de treinamento com centenas de bilhões de parâmetros e otimizações de desempenho que reduzem o tempo de treinamento do modelo e custo em até 20%.

Para saber mais sobre a biblioteca paralela do modelo SageMaker, consulte Biblioteca de paralelismo de modelo SageMaker v2 documentação. Você também pode consultar nosso cadernos de exemplo para começar.

Novos recursos que simplificam e aceleram o treinamento de modelos grandes

Esta postagem discute os recursos mais recentes incluídos na versão v2.0 da biblioteca paralela de modelos SageMaker. Esses recursos melhoram a usabilidade da biblioteca, expandem a funcionalidade e aceleram o treinamento. Nas seções a seguir, resumimos os novos recursos e discutimos como você pode usar a biblioteca para acelerar o treinamento de modelos grandes.

Alinhando SMP com PyTorch de código aberto

Desde o seu lançamento em 2020, o SMP permitiu treinamento em grande escala e alto desempenho em instâncias de computação do SageMaker. Com esta versão principal mais recente do SMP, a biblioteca simplifica a experiência do usuário alinhando suas APIs com PyTorch de código aberto.

Ofertas PyTorch Paralelismo de dados totalmente fragmentado (FSDP) como seu principal método para suportar grandes cargas de trabalho de treinamento em muitos dispositivos de computação. Conforme demonstrado no trecho de código a seguir, as APIs atualizadas do SMP para técnicas como paralelismo de dados fragmentados refletem as do PyTorch. Você pode simplesmente executar import torch.sagemaker e use-o no lugar de torch.

## training_script.py
import torch.sagemaker as tsm
tsm.init()

# Set up a PyTorch model
model = ...

# Wrap the PyTorch model using the PyTorch FSDP module
model = FSDP(
model,
...
)

optimizer = ...
...

Com essas atualizações nas APIs do SMP, agora você pode obter os benefícios de desempenho do SageMaker e da biblioteca SMP sem revisar seus scripts de treinamento PyTorch FSDP existentes. Esse paradigma também permite que você use a mesma base de código ao treinar no local e no SageMaker, simplificando a experiência do usuário para clientes que treinam em vários ambientes.

Para obter mais informações sobre como habilitar o SMP com seus scripts de treinamento PyTorch FSDP existentes, consulte Comece com o SMP.

Integrando o paralelismo de tensores para permitir o treinamento em clusters massivos

Esta versão do SMP também expande os recursos do PyTorch FSDP para incluir técnicas de paralelismo de tensores. Um problema ao usar apenas o paralelismo de dados fragmentados é que você pode encontrar problemas de convergência à medida que aumenta o tamanho do cluster. Isso ocorre porque a fragmentação de parâmetros, gradientes e o estado do otimizador em classificações paralelas de dados também aumenta o tamanho global do lote; em grandes clusters, esse tamanho de lote global pode ser empurrado para além do limite abaixo do qual o modelo convergiria. Você precisa incorporar uma técnica de paralelismo adicional que não exija um aumento no tamanho global do lote à medida que você dimensiona seu cluster.

Para mitigar esse problema, o SMP v2.0 introduz a capacidade de compor o paralelismo de dados fragmentados com paralelismo de tensor. O paralelismo tensorial permite que o tamanho do cluster aumente sem alterar o tamanho global do lote ou afetar a convergência do modelo. Com esse recurso, você pode aumentar com segurança o rendimento do treinamento provisionando clusters com 256 nós ou mais.

Hoje, o paralelismo de tensores com PyTorch FSDP está disponível apenas com SMP v2. O SMP v2 permite ativar essa técnica com algumas linhas de alteração de código e desbloquear treinamento estável mesmo em clusters grandes. SMP v2 integra-se com Motor Transformador por sua implementação de paralelismo de tensor e o torna compatível com as APIs PyTorch FSDP. Você pode habilitar o paralelismo de tensor PyTorch FSDP e SMP simultaneamente sem fazer nenhuma alteração em seu modelo PyTorch ou na configuração PyTorch FSDP. Os trechos de código a seguir mostram como configurar o dicionário de configuração SMP no formato JSON e adicionar o módulo de inicialização SMP torch.sagemaker.init(), que aceita o dicionário de configuração no back-end quando você inicia o trabalho de treinamento, para seu script de treinamento.

A configuração do SMP é a seguinte:

{
"tensor_parallel_degree": 8,
"tensor_parallel_seed": 0
}

No seu script de treinamento, use o seguinte código:

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model)

Para saber mais sobre o uso do paralelismo tensorial no SMP, consulte o paralelismo tensorial seção de nossa documentação.

Use recursos avançados para acelerar o treinamento de modelos em até 20%

Além de permitir treinamento distribuído em clusters com centenas de instâncias, o SMP também oferece técnicas de otimização que podem acelerar o treinamento de modelos em até 20%. Nesta seção, destacamos algumas dessas otimizações. Para saber mais, consulte o recursos principais seção de nossa documentação.

Fragmentação híbrida

O paralelismo de dados fragmentados é uma técnica de treinamento distribuído que economiza memória e que divide o estado de um modelo (parâmetros do modelo, gradientes e estados do otimizador) entre dispositivos. Esse menor consumo de memória permite que você ajuste um modelo maior em seu cluster ou aumente o tamanho do lote. No entanto, o paralelismo de dados fragmentados também aumenta os requisitos de comunicação do seu trabalho de formação porque os artefactos do modelo fragmentado são frequentemente recolhidos de diferentes dispositivos durante o treino. Dessa forma, o grau de fragmentação é uma configuração importante que compensa o consumo de memória e a sobrecarga de comunicação.

Por padrão, os fragmentos PyTorch FSDP modelam artefatos em todos os dispositivos aceleradores em seu cluster. Dependendo do seu trabalho de treinamento, esse método de fragmentação pode aumentar a sobrecarga de comunicação e criar um gargalo. Para ajudar com isso, a biblioteca SMP oferece paralelismo de dados fragmentados híbridos configuráveis ​​​​sobre PyTorch FSDP. Esse recurso permite definir o grau de fragmentação ideal para sua carga de trabalho de treinamento. Basta especificar o grau de fragmentação em um objeto JSON de configuração e incluí-lo em seu script de treinamento SMP.

A configuração do SMP é a seguinte:

{ "hybrid_shard_degree": 16 }

Para saber mais sobre as vantagens do paralelismo de dados fragmentados híbridos, consulte Dimensionamento quase linear de treinamento de modelo gigantesco na AWS. Para obter mais informações sobre a implementação de fragmentação híbrida com seu script de treinamento FSDP existente, consulte paralelismo de dados compartilhados híbridos em nossa documentação.

Use as operações de comunicação coletiva SMDDP otimizadas para a infraestrutura da AWS

Você pode usar a biblioteca SMP com o Biblioteca de paralelismo de dados distribuídos SageMaker (SMDDP) para acelerar suas cargas de trabalho de treinamento distribuídas. SMDDP inclui um otimizado AllGather operação de comunicação coletiva projetada para melhor desempenho em instâncias aceleradas p4d e p4de do SageMaker. No treinamento distribuído, as operações de comunicação coletiva são usadas para sincronizar informações entre os trabalhadores da GPU. AllGather é uma das principais operações de comunicação coletiva normalmente usadas no paralelismo de dados fragmentados para materializar os parâmetros da camada antes das etapas de computação direta e reversa. Para trabalhos de formação que são dificultados pela comunicação, operações colectivas mais rápidas podem reduzir o tempo e o custo da formação, sem efeitos secundários na convergência.

Para usar a biblioteca SMDDP, você só precisa adicionar duas linhas de código ao seu script de treinamento:

import torch.distributed as dist

# Initialize with SMDDP
import smdistributed.dataparallel.torch.torch_smddp
dist.init_process_group(backend="smddp") # Replacing "nccl"

# Initialize with SMP
import torch.sagemaker as tsm
tsm.init()

Além do SMP, o SMDDP suporta PyTorch FSDP e DeepSpeed ​​de código aberto. Para saber mais sobre a biblioteca SMDDP, consulte Execute treinamento distribuído com a biblioteca de paralelismo de dados distribuídos SageMaker.

Descarregamento de ativação

Normalmente, a passagem para frente do treinamento do modelo calcula ativações em cada camada e as mantém na memória da GPU até que a passagem para trás da camada correspondente termine. Essas ativações armazenadas podem consumir uma quantidade significativa de memória da GPU durante o treinamento. O descarregamento de ativação é uma técnica que, em vez disso, move esses tensores para a memória da CPU após a passagem direta e depois os recupera de volta para a GPU quando são necessários. Essa abordagem pode reduzir substancialmente o uso de memória da GPU durante o treinamento.

Embora o PyTorch suporte o descarregamento de ativação, sua implementação é ineficiente e pode fazer com que as GPUs fiquem ociosas enquanto as ativações são recuperadas da CPU durante uma passagem para trás. Isso pode causar degradação significativa do desempenho ao usar o descarregamento de ativação.

O SMP v2 oferece um algoritmo de descarregamento de ativação otimizado que pode melhorar o desempenho do treinamento. A implementação do SMP pré-busca ativações antes que elas sejam necessárias na GPU, reduzindo o tempo ocioso.

Como o SMP é construído com base nas APIs do PyTorch, permitir o descarregamento de ativação otimizado requer apenas algumas linhas de alteração de código. Basta adicionar as configurações associadas (sm_activation_offloading e activation_loading_horizon parâmetros) e inclua-os em seu roteiro de treinamento.

A configuração do SMP é a seguinte:

{
"activation_loading_horizon": 2,
"sm_activation_offloading": True
}

No script de treinamento, use o seguinte código:

import torch.sagemaker as tsm
tsm.init()

# Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
apply_activation_checkpointing,
offload_wrapper,
)

model = FSDP(...)

# Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
model,
check_fn=checkpoint_tformer_layers_policy,
)

model = offload_wrapper(model)

Para saber mais sobre as ferramentas de ponto de verificação PyTorch de código aberto para descarregamento de ativação, consulte o checkpoint_wrapper.py script no repositório PyTorch GitHub e Ponto de verificação de ativação na postagem do blog PyTorch Dimensionando modelos de base multimodal no TorchMultimodal com Pytorch Distributed. Para saber mais sobre a implementação otimizada de descarregamento de ativação do SMP, consulte o descarga de ativação seção de nossa documentação.

Além da fragmentação híbrida, do SMDDP e do descarregamento de ativação, o SMP oferece otimizações adicionais que podem acelerar sua grande carga de trabalho de treinamento de modelos. Isso inclui pontos de verificação de ativação otimizados, inicialização atrasada de parâmetros e outros. Para saber mais, consulte o recursos principais seção de nossa documentação.

Conclusão

À medida que os conjuntos de dados, tamanhos de modelos e clusters de treinamento continuam a crescer, o treinamento distribuído eficiente é cada vez mais crítico para a entrega oportuna e acessível de modelos e produtos. A versão mais recente da biblioteca paralela de modelo SageMaker ajuda você a conseguir isso, reduzindo alterações de código e alinhando-se com APIs PyTorch FSDP, permitindo o treinamento em clusters massivos por meio de paralelismo de tensor e otimizações que podem reduzir o tempo de treinamento em até 20%.

Para começar a usar o SMP v2, consulte nosso documentação e nosso cadernos de amostra.

Sobre os autores

Robert Van Dusen é gerente de produto sênior da Amazon SageMaker. Ele lidera estruturas, compiladores e técnicas de otimização para treinamento de aprendizado profundo.

Luis Quintela é o Software Developer Manager para a biblioteca paralela de modelos do AWS SageMaker. Em seu tempo livre, ele pode ser encontrado andando de Harley na área da baía de SF.

Gautama Kumar é engenheiro de software com AWS AI Deep Learning. Ele é apaixonado por construir ferramentas e sistemas para IA. Nas horas vagas, ele gosta de andar de bicicleta e ler livros.

Rahul Huilgol é Engenheiro Sênior de Desenvolvimento de Software em Distributed Deep Learning na Amazon Web Services.

Carimbo de hora:

Mais de Aprendizado de máquina da AWS