Treinamento distribuído com Amazon EKS e Torch Distributed Elastic

O treinamento distribuído de modelos de aprendizado profundo está se tornando cada vez mais importante à medida que os tamanhos de dados estão crescendo em muitos setores. Muitos aplicativos em visão computacional e processamento de linguagem natural agora exigem treinamento de modelos de aprendizado profundo, que estão crescendo exponencialmente em complexidade e geralmente são treinados com centenas de terabytes de dados. Torna-se então importante usar uma vasta infraestrutura em nuvem para dimensionar o treinamento de modelos tão grandes.

Os desenvolvedores podem usar estruturas de código aberto, como o PyTorch, para projetar facilmente arquiteturas de modelo intuitivas. No entanto, dimensionar o treinamento desses modelos em vários nós pode ser um desafio devido à maior complexidade da orquestração.

O treinamento de modelo distribuído consiste principalmente em dois paradigmas:

• Modelo paralelo – No treinamento paralelo de modelo, o próprio modelo é tão grande que não cabe na memória de uma única GPU, e várias GPUs são necessárias para treinar o modelo. O modelo GPT-3 da Open AI com 175 bilhões de parâmetros treináveis ​​(aproximadamente 350 GB de tamanho) é um bom exemplo disso.
• Dados paralelos – No treinamento paralelo de dados, o modelo pode residir em uma única GPU, mas, como os dados são muito grandes, pode levar dias ou semanas para treinar um modelo. A distribuição dos dados em vários nós de GPU pode reduzir significativamente o tempo de treinamento.

Neste post, fornecemos uma arquitetura de exemplo para treinar modelos PyTorch usando o Elástico distribuído de tocha framework de forma paralela de dados distribuídos usando Serviço Amazon Elastic Kubernetes (Amazon EKS).

Pré-requisitos

Para replicar os resultados relatados nesta postagem, o único pré-requisito é uma conta da AWS. Nesta conta, criamos um cluster EKS e um Amazon FSx para Lustre sistema de arquivo. Também enviamos imagens de contêiner para um Registro do Amazon Elastic Container (Amazon ECR) na conta. As instruções para configurar esses componentes são fornecidas conforme necessário em todo o post.

Agrupamentos EKS

O Amazon EKS é um serviço de contêiner gerenciado para executar e dimensionar aplicativos Kubernetes na AWS. Com o Amazon EKS, você pode executar tarefas de treinamento distribuídas com eficiência usando as últimas Amazon Elastic Compute Nuvem (Amazon EC2) sem precisar instalar, operar e manter seu próprio plano de controle ou nós. É um popular orquestrador para fluxos de trabalho de aprendizado de máquina (ML) e IA. Um cluster EKS típico na AWS se parece com a figura a seguir.

Lançamos um projeto de código aberto, AWS DevOps para EKS (aws-do-eks), que fornece uma grande coleção de scripts e ferramentas fáceis de usar e configuráveis ​​para provisionar clusters EKS e executar trabalhos de treinamento distribuídos. Este projeto é construído seguindo os princípios da Fazer estrutura: Simplicidade, Flexibilidade e Universalidade. Você pode configurar o cluster desejado usando o eks.conf arquivo e, em seguida, inicie-o executando o eks-create.sh roteiro. Instruções detalhadas são fornecidas no GitHub repo.

Treinar modelos PyTorch usando Torch Distributed Elastic

O Torch Distributed Elastic (TDE) é uma biblioteca PyTorch nativa para treinamento de modelos de aprendizado profundo em grande escala, onde é fundamental dimensionar recursos de computação dinamicamente com base na disponibilidade. o TorchElastic Controller para Kubernetes é uma implementação nativa do Kubernetes para TDE que gerencia automaticamente o ciclo de vida dos pods e serviços necessários para o treinamento do TDE. Ele permite dimensionar dinamicamente os recursos de computação durante o treinamento, conforme necessário. Ele também fornece treinamento tolerante a falhas recuperando trabalhos de falha de nó.

Neste post, discutimos as etapas para treinar o PyTorch EficienteNet-B7 e ResNet50 modelos usando IMAGEnet dados de forma distribuída com TDE. Usamos o PyTorch DistribuídoDataParallel API e o controlador Kubernetes TorchElastic e execute nossos trabalhos de treinamento em um cluster EKS contendo vários nós de GPU. O diagrama a seguir mostra o diagrama de arquitetura para este treinamento de modelo.

O TorchElastic para Kubernetes consiste principalmente em dois componentes: o TorchElastic Kubernetes Controller (TEC) e o servidor de parâmetros (etcd). O controlador é responsável por monitorar e gerenciar as tarefas de treinamento e o servidor de parâmetros acompanha os nós do trabalhador para sincronização distribuída e descoberta de pares.

Para que os pods de treinamento acessem os dados, precisamos de um volume de dados compartilhado que possa ser montado por cada pod. Algumas opções para volumes compartilhados por meio de Interface de armazenamento de contêiner (CSI) drivers incluídos no AWS DevOps para EKS e guarante que os mesmos estão Sistema de arquivos elástico da Amazon (Amazônia EFS) e FSx para brilho.

Configuração do cluster

Em nossa configuração de cluster, usamos uma instância c5.2xlarge para pods do sistema. Usamos três instâncias p4d.24xlarge como pods de trabalho para treinar um modelo EfficientNet. Para o treinamento ResNet50, usamos instâncias p3.8xlarge como pods de trabalho. Além disso, usamos um sistema de arquivos compartilhado FSx para armazenar nossos dados de treinamento e artefatos de modelo.

As instâncias AWS p4d.24xlarge são equipadas com Adaptador de tecido elástico (EFA) para fornecer rede entre nós. Discutimos a EFA mais adiante no post. Para habilitar a comunicação por meio do EFA, precisamos configurar a configuração do cluster por meio de um arquivo .yaml. Um arquivo de exemplo é fornecido no repositório do GitHub.

Depois que esse arquivo .yaml estiver configurado corretamente, podemos iniciar o cluster usando o script fornecido no repositório do GitHub:

./eks-create.sh

Consulte o GitHub repo para instruções detalhadas.

Praticamente não há diferença entre executar trabalhos em p4d.24xlarge e p3.8xlarge. As etapas descritas neste post funcionam para ambos. A única diferença é a disponibilidade do EFA em instâncias p4d.24xlarge. Para modelos menores como ResNet50, a rede padrão comparada à rede EFA tem um impacto mínimo na velocidade do treinamento.

FSx para sistema de arquivos Lustre

O FSx foi projetado para cargas de trabalho de computação de alto desempenho e fornece latência de menos de um milissegundo usando volumes de armazenamento de unidade de estado sólido. Escolhemos o FSx porque ele forneceu melhor desempenho à medida que escalamos para um grande número de nós. Um detalhe importante a ser observado é que o FSx só pode existir em uma única zona de disponibilidade. Portanto, todos os nós que acessam o sistema de arquivos FSx devem existir na mesma zona de disponibilidade que o sistema de arquivos FSx. Uma maneira de fazer isso é especificar a zona de disponibilidade relevante no arquivo .yaml do cluster para os grupos de nós específicos antes de criar o cluster. Como alternativa, podemos modificar a parte de rede do grupo de dimensionamento automático para esses nós após a configuração do cluster e limitá-lo ao uso de uma única sub-rede. Isso pode ser feito facilmente no console do Amazon EC2.

Supondo que o cluster EKS esteja funcionando e o ID da sub-rede para a zona de disponibilidade seja conhecido, podemos configurar um sistema de arquivos FSx fornecendo as informações necessárias no fsx.conf arquivo conforme descrito no readme e executando o implantar.sh script no fsx pasta. Isso configura a política e o grupo de segurança corretos para acessar o sistema de arquivos. O script também instala o Driver CSI para FSx como um daemonset. Por fim, podemos criar a declaração de volume persistente FSx no Kubernetes aplicando um único arquivo .yaml:

kubectl apply -f fsx-pvc-dynamic.yaml

Isso cria um sistema de arquivos FSx na zona de disponibilidade especificada no fsx.conf arquivo e também cria uma declaração de volume persistente fsx-pvc, que pode ser montado por qualquer um dos pods no cluster em um modo de leitura-gravação-muitos (RWX).

Em nosso experimento, usamos dados completos do ImageNet, que contêm mais de 12 milhões de imagens de treinamento divididas em 1,000 classes. Os dados podem ser baixados do Site da ImageNet. A bola TAR original tem vários diretórios, mas para nosso treinamento de modelo, estamos interessados ​​apenas em ILSVRC/Data/CLS-LOC/, que inclui o train e val subdiretórios. Antes do treinamento, precisamos reorganizar as imagens no val subdiretório para corresponder à estrutura de diretório exigida pelo PyTorch Pasta de imagens classe. Isso pode ser feito usando um simples Script Python após os dados serem copiados para o volume persistente na próxima etapa.

Para copiar os dados de um Serviço de armazenamento simples da Amazon (Amazon S3) para o sistema de arquivos FSx, criamos uma imagem do Docker que inclui scripts para essa tarefa. Um exemplo Dockerfile e um script de shell estão incluídos no csi pasta dentro do repositório do GitHub. Podemos construir a imagem usando o build.sh script e, em seguida, envie-o para o Amazon ECR usando o push.sh roteiro. Antes de usar esses scripts, precisamos fornecer o URI correto para o repositório ECR no .env arquivo na pasta raiz do repositório GitHub. Depois de enviarmos a imagem do Docker para o Amazon ECR, podemos iniciar um pod para copiar os dados aplicando o arquivo .yaml relevante:

kubectl apply -f fsx-data-prep-pod.yaml

O pod executa automaticamente o script preparação de dados.sh para copiar os dados do Amazon S3 para o volume compartilhado. Como os dados do ImageNet têm mais de 12 milhões de arquivos, o processo de cópia leva algumas horas. O script Python imagenet_data_prep.py também é executado para reorganizar o val conjunto de dados como esperado pelo PyTorch.

Aceleração da rede

Podemos usar o Elastic Fabric Adapter (EFA) em combinação com tipos de instância EC2 compatíveis para acelerar o tráfego de rede entre os nós GPU em seu cluster. Isso pode ser útil ao executar grandes trabalhos de treinamento distribuídos em que a comunicação de rede padrão pode ser um gargalo. Os scripts para implantar e testar o plug-in do dispositivo EFA no cluster EKS que usamos aqui estão incluídos no efa-device-plugin pasta no repositório do GitHub. Para habilitar um trabalho com EFA em seu cluster EKS, além de os nós do cluster terem o hardware e software necessários, o plug-in do dispositivo EFA precisa ser implantado no cluster e seu contêiner de trabalho precisa ter CUDA e NCCL compatíveis versões instalado.

Para demonstrar a execução de testes NCCL e a avaliação do desempenho do EFA em instâncias p4d.24xlarge, primeiro devemos implantar o operador Kubeflow MPI executando o correspondente implantar.sh script no operador mpi pasta. Em seguida, executamos o implantar.sh script e atualize o test-efa-nccl.yaml manifeste assim limites e solicitações de recursos vpc.amazonaws.com são definidos como 4. Os quatro adaptadores EFA disponíveis nos nós p4d.24xlarge são agrupados para fornecer o máximo rendimento.

Execute kubectl apply -f ./test-efa-nccl.yaml para aplicar o teste e, em seguida, exibir os logs do pod de teste. A linha a seguir na saída do log confirma que o EFA está sendo usado:

NCCL INFO NET/OFI Selected Provider is efa

Os resultados do teste devem ser semelhantes à seguinte saída:

[1,0]<stdout>:#                                                       out-of-place                       in-place
[1,0]<stdout>:#       size         count      type   redop     time   algbw   busbw  error     time   algbw   busbw  error
[1,0]<stdout>:#        (B)    (elements)                       (us)  (GB/s)  (GB/s)            (us)  (GB/s)  (GB/s)
[1,0]<stdout>:           8             2     float     sum    629.7    0.00    0.00  2e-07    631.4    0.00    0.00  1e-07
[1,0]<stdout>:          16             4     float     sum    630.5    0.00    0.00  1e-07    628.1    0.00    0.00  1e-07
[1,0]<stdout>:          32             8     float     sum    627.6    0.00    0.00  1e-07    628.2    0.00    0.00  1e-07
[1,0]<stdout>:          64            16     float     sum    633.6    0.00    0.00  1e-07    628.4    0.00    0.00  6e-08
[1,0]<stdout>:         128            32     float     sum    627.5    0.00    0.00  6e-08    632.0    0.00    0.00  6e-08
[1,0]<stdout>:         256            64     float     sum    634.5    0.00    0.00  6e-08    636.5    0.00    0.00  6e-08
[1,0]<stdout>:         512           128     float     sum    634.8    0.00    0.00  6e-08    635.2    0.00    0.00  6e-08
[1,0]<stdout>:        1024           256     float     sum    646.6    0.00    0.00  2e-07    643.6    0.00    0.00  2e-07
[1,0]<stdout>:        2048           512     float     sum    745.0    0.00    0.01  5e-07    746.0    0.00    0.01  5e-07
[1,0]<stdout>:        4096          1024     float     sum    958.2    0.00    0.01  5e-07    955.8    0.00    0.01  5e-07
[1,0]<stdout>:        8192          2048     float     sum    963.0    0.01    0.02  5e-07    954.5    0.01    0.02  5e-07
[1,0]<stdout>:       16384          4096     float     sum    955.0    0.02    0.03  5e-07    955.5    0.02    0.03  5e-07
[1,0]<stdout>:       32768          8192     float     sum    975.5    0.03    0.06  5e-07   1009.0    0.03    0.06  5e-07
[1,0]<stdout>:       65536         16384     float     sum   1353.4    0.05    0.09  5e-07   1343.5    0.05    0.09  5e-07
[1,0]<stdout>:      131072         32768     float     sum   1395.9    0.09    0.18  5e-07   1392.6    0.09    0.18  5e-07
[1,0]<stdout>:      262144         65536     float     sum   1476.7    0.18    0.33  5e-07   1536.3    0.17    0.32  5e-07
[1,0]<stdout>:      524288        131072     float     sum   1560.3    0.34    0.63  5e-07   1568.3    0.33    0.63  5e-07
[1,0]<stdout>:     1048576        262144     float     sum   1599.2    0.66    1.23  5e-07   1595.3    0.66    1.23  5e-07
[1,0]<stdout>:     2097152        524288     float     sum   1671.1    1.25    2.35  5e-07   1672.5    1.25    2.35  5e-07
[1,0]<stdout>:     4194304       1048576     float     sum   1785.1    2.35    4.41  5e-07   1780.3    2.36    4.42  5e-07
[1,0]<stdout>:     8388608       2097152     float     sum   2133.6    3.93    7.37  5e-07   2135.0    3.93    7.37  5e-07
[1,0]<stdout>:    16777216       4194304     float     sum   2650.9    6.33   11.87  5e-07   2649.9    6.33   11.87  5e-07
[1,0]<stdout>:    33554432       8388608     float     sum   3422.0    9.81   18.39  5e-07   3478.7    9.65   18.09  5e-07
[1,0]<stdout>:    67108864      16777216     float     sum   4783.2   14.03   26.31  5e-07   4782.6   14.03   26.31  5e-07
[1,0]<stdout>:   134217728      33554432     float     sum   7216.9   18.60   34.87  5e-07   7240.9   18.54   34.75  5e-07
[1,0]<stdout>:   268435456      67108864     float     sum    12738   21.07   39.51  5e-07    12802   20.97   39.31  5e-07
[1,0]<stdout>:   536870912     134217728     float     sum    24375   22.03   41.30  5e-07    24403   22.00   41.25  5e-07
[1,0]<stdout>:  1073741824     268435456     float     sum    47904   22.41   42.03  5e-07    47893   22.42   42.04  5e-07
[1,4]<stdout>:test-efa-nccl-worker-0:33:33 [4] NCCL INFO comm 0x7fd4a0000f60 rank 4 nranks 16 cudaDev 4 busId 901c0 - Destroy COMPLETE
[1,0]<stdout>:# Out of bounds values : 0 OK
[1,0]<stdout>:# Avg bus bandwidth    : 8.23785

Podemos observar nos resultados do teste que a taxa de transferência máxima é de cerca de 42 GB/s e a largura de banda média do barramento é de aproximadamente 8 GB.

Também realizamos experimentos com um único adaptador EFA habilitado e sem adaptadores EFA. Todos os resultados estão resumidos na tabela a seguir.

Também descobrimos que para modelos relativamente pequenos como o ImageNet, o uso de rede acelerada reduz o tempo de treinamento por época apenas com 5 a 8% no tamanho de lote de 64. Para modelos maiores e tamanhos de lote menores, quando é necessária maior comunicação de rede de pesos , o uso de rede acelerada tem maior impacto. Observamos uma diminuição do tempo de treinamento de época com 15–18% para o treinamento do EfficientNet-B7 com tamanho de lote 1. O impacto real do EFA em seu treinamento dependerá do tamanho do seu modelo.

Monitoramento de GPU

Antes de executar o trabalho de treinamento, também podemos configurar Amazon CloudWatch métricas para visualizar a utilização da GPU durante o treinamento. Pode ser útil saber se os recursos estão sendo usados ​​de maneira ideal ou identificar potencialmente a falta de recursos e os gargalos no processo de treinamento.

Os scripts relevantes para configurar o CloudWatch estão localizados no métricas de GPU pasta. Primeiro, criamos uma imagem do Docker com amazon-cloudwatch-agent e nvidia-smi. Podemos usar o Dockerfile no gpu-metrics pasta para criar esta imagem. Supondo que o registro ECR já esteja definido no .env arquivo da etapa anterior, podemos construir e enviar a imagem usando build.sh e push.sh. Após isso, executar o deploy.sh script conclui automaticamente a configuração. Ele lança um daemonset com amazon-cloudwatch-agent e envia várias métricas para o CloudWatch. As métricas da GPU aparecem sob o CWAgent namespace no console do CloudWatch. O restante das métricas do cluster são exibidos sob o ContainerInsights espaço para nome.

Treinamento de modelo

Todos os scripts necessários para o treinamento do PyTorch estão localizados no elástico pasta no repositório do GitHub. Antes de iniciar o job de treinamento, precisamos executar o etcd servidor, que é usado pelo TEC para descoberta do trabalhador e troca de parâmetros. o implantar.sh script no elasticjob pasta faz exatamente isso.

Para aproveitar o EFA em instâncias p4d.24xlarge, precisamos usar uma imagem específica do Docker disponível no Galeria pública do Amazon ECR que suporta a comunicação NCCL através da EFA. Só precisamos copiar nosso código de treinamento para esta imagem do Docker. o dockerfile sob a amostras A pasta cria uma imagem a ser usada ao executar o trabalho de treinamento em instâncias p4d. Como sempre, podemos usar o construir.sh e empurrar.sh scripts na pasta para construir e enviar a imagem.

A imagenet-efa.yaml arquivo descreve o trabalho de treinamento. Esse arquivo .yaml configura os recursos necessários para executar o trabalho de treinamento e também monta o volume persistente com os dados de treinamento configurados na seção anterior.

Algumas coisas merecem destaque aqui. O número de réplicas deve ser definido para o número de nós disponíveis no cluster. No nosso caso, definimos como 3 porque tínhamos três nós p4d.24xlarge. No imagenet-efa.yaml arquivo, o nvidia.com/gpu parâmetro em recursos e nproc_per_node para args deve ser definido para o número de GPUs por nó, que no caso de p4d.24xlarge é 8. Além disso, o argumento do trabalhador para o script Python define o número de CPUs por processo. Escolhemos o valor 4 porque, em nossos experimentos, isso fornece desempenho ideal ao executar em instâncias p4d.24xlarge. Essas configurações são necessárias para maximizar o uso de todos os recursos de hardware disponíveis no cluster.

Quando o trabalho está em execução, podemos observar o uso da GPU no CloudWatch para todas as GPUs no cluster. Veja a seguir um exemplo de um de nossos trabalhos de treinamento com três nós p4d.24xlarge no cluster. Aqui selecionamos uma GPU de cada nó. Com as configurações mencionadas anteriormente, o uso da GPU está próximo de 100% durante a fase de treinamento da época para todos os nós do cluster.

Para treinar um modelo ResNet50 usando instâncias p3.8xlarge, precisamos exatamente das mesmas etapas descritas para o treinamento EfficientNet usando p4d.24xlarge. Também podemos usar a mesma imagem do Docker. Conforme mencionado anteriormente, as instâncias p3.8xlarge não são equipadas com EFA. No entanto, para o modelo ResNet50, isso não é uma desvantagem significativa. o imagenet-fsx.yaml O script fornecido no repositório do GitHub configura o trabalho de treinamento com recursos apropriados para o tipo de nó p3.8xlarge. A tarefa usa o mesmo conjunto de dados do sistema de arquivos FSx.

Dimensionamento de GPU

Executamos alguns experimentos para observar como o tempo de treinamento é dimensionado para o modelo EfficientNet-B7 aumentando o número de GPUs. Para fazer isso, alteramos o número de réplicas de 1 para 3 em nosso arquivo .yaml de treinamento para cada execução de treinamento. Observamos apenas o tempo para uma única época ao usar o conjunto de dados completo do ImageNet. A figura a seguir mostra os resultados de nosso experimento de dimensionamento de GPU. A linha pontilhada vermelha representa como o tempo de treinamento deve diminuir de uma execução usando 8 GPUs aumentando o número de GPUs. Como podemos ver, o dimensionamento está bem próximo do esperado.

Da mesma forma, obtivemos o gráfico de dimensionamento de GPU para treinamento ResNet50 em instâncias p3.8xlarge. Para este caso, alteramos as réplicas em nosso arquivo .yaml de 1 para 4. Os resultados desse experimento são mostrados na figura a seguir.

limpar

É importante reduzir os recursos após o treinamento do modelo para evitar custos associados à execução de instâncias ociosas. Com cada script que cria recursos, o GitHub repo fornece um script correspondente para excluí-los. Para limpar nossa configuração, devemos excluir o sistema de arquivos FSx antes de excluir o cluster porque ele está associado a uma sub-rede na VPC do cluster. Para excluir o sistema de arquivos FSx, basta executar o seguinte comando (de dentro do fsx pasta):

kubectl delete -f fsx-pvc-dynamic.yaml
./delete.sh

Observe que isso não apenas excluirá o volume persistente, mas também excluirá o sistema de arquivos FSx e todos os dados no sistema de arquivos serão perdidos. Quando esta etapa for concluída, podemos excluir o cluster usando o script a seguir no ex pasta:

./eks-delete.sh

Isso excluirá todos os pods existentes, removerá o cluster e excluirá a VPC criada no início.

Conclusão

Nesta postagem, detalhamos as etapas necessárias para executar o treinamento de modelo paralelo de dados distribuídos do PyTorch em clusters EKS. Essa tarefa pode parecer assustadora, mas o AWS DevOps para EKS O projeto criado pela equipe de ML Frameworks da AWS fornece todos os scripts e ferramentas necessários para simplificar o processo e tornar o treinamento de modelo distribuído facilmente acessível.

Para mais informações sobre as tecnologias usadas neste post, visite Amazon EX e Elástico distribuído de tocha. Recomendamos que você aplique a abordagem descrita aqui aos seus próprios casos de uso de treinamento distribuído.

Recursos

Sobre os autores

Imran Younus é um arquiteto de soluções principal para a equipe de estruturas de ML na AWS. Ele se concentra em cargas de trabalho de aprendizado de máquina e aprendizado profundo em grande escala em serviços da AWS, como Amazon EKS e AWS ParallelCluster. Possui vasta experiência em aplicações de Deep Leaning em Visão Computacional e IoT Industrial. Imran obteve seu PhD em Física de Partículas de Alta Energia, onde esteve envolvido na análise de dados experimentais em escalas de peta-byte.

Alex Iankoulski é um arquiteto de infraestrutura e software full-stack que gosta de fazer um trabalho profundo e prático. Atualmente, ele é o principal arquiteto de soluções para aprendizado de máquina autogerenciado na AWS. Em sua função, ele se concentra em ajudar os clientes com conteinerização e orquestração de cargas de trabalho de ML e IA em serviços da AWS baseados em contêiner. Ele também é o autor do código aberto Fazer estrutura e um capitão do Docker que adora aplicar tecnologias de contêineres para acelerar o ritmo da inovação enquanto resolve os maiores desafios do mundo. Nos últimos 10 anos, Alex trabalhou no combate às mudanças climáticas, democratizando a IA e o ML, tornando as viagens mais seguras, a saúde melhor e a energia mais inteligente.