Execute o PyTorch Lightning e o PyTorch DDP nativo no Amazon SageMaker Training, com o Amazon Search

Tantos dados, tão pouco tempo. Especialistas em aprendizado de máquina (ML), cientistas de dados, engenheiros e entusiastas encontraram esse problema em todo o mundo. Do processamento de linguagem natural à visão computacional, tabular às séries temporais e tudo mais, o antigo problema de otimizar a velocidade ao executar dados no maior número possível de GPUs inspirou inúmeras soluções. Hoje, temos o prazer de anunciar recursos para desenvolvedores PyTorch usando estruturas nativas de código aberto, como Relâmpago PyTorch e PyTorch DDP, que simplificará seu caminho para a nuvem.

Amazon Sage Maker é um serviço totalmente gerenciado para ML, e o treinamento de modelo SageMaker é um ambiente de computação otimizado para treinamento de alto desempenho em escala. O treinamento de modelos do SageMaker oferece uma experiência de treinamento remoto com um plano de controle contínuo para treinar e reproduzir facilmente modelos de ML com alto desempenho e baixo custo. Temos o prazer de anunciar novos recursos no portfólio de treinamento do SageMaker que tornam a execução do PyTorch em escala ainda mais fácil e acessível:

1. O PyTorch Lightning agora pode ser integrado a Biblioteca paralela de dados distribuídos do SageMaker com apenas uma linha de mudança de código.
2. Treinamento do modelo SageMaker agora tem suporte para PyTorch Distributed Data Parallel nativo com backend NCCL, permitindo que os desenvolvedores migrem para o SageMaker com mais facilidade do que nunca.

Nesta postagem, discutimos esses novos recursos e também aprendemos como o Amazon Search executou o PyTorch Lightning com o back-end de treinamento distribuído otimizado no SageMaker para acelerar o tempo de treinamento do modelo.

Antes de mergulhar no estudo de caso do Amazon Search, para aqueles que não estão familiarizados, gostaríamos de fornecer algumas informações sobre Biblioteca paralela de dados distribuídos do SageMaker. Em 2020, desenvolvemos e lançamos uma configuração de cluster personalizada para descida de gradiente distribuída em escala que aumenta a eficiência geral do cluster, apresentada na Amazon Science como Arenque. Usando o melhor dos servidores de parâmetros e topologias baseadas em anel, o SageMaker Distributed Data Parallel (SMDDP) é otimizado para o Amazon Elastic Compute Nuvem (Amazon EC2) topologia de rede, incluindo EFA. Para tamanhos de cluster maiores, o SMDDP é capaz de fornecer melhorias de taxa de transferência de 20 a 40% em relação ao Horovod (TensorFlow) e ao paralelo de dados distribuídos do PyTorch. Para tamanhos de cluster menores e modelos compatíveis, recomendamos o Compilador de treinamento SageMaker, que é capaz de diminuir o tempo total de trabalho em até 50%.

Destaque do cliente: PyTorch Lightning no back-end otimizado do SageMaker com o Amazon Search

O Amazon Search é responsável pela experiência de pesquisa e descoberta em Amazon.com. Ele potencializa a experiência de pesquisa para clientes que procuram produtos para comprar na Amazon. Em alto nível, o Amazon Search cria um índice para todos os produtos vendidos na Amazon.com. Quando um cliente insere uma consulta, o Amazon Search usa várias técnicas de ML, incluindo modelos de aprendizado profundo, para corresponder produtos relevantes e interessantes à consulta do cliente. Em seguida, ele classifica os produtos antes de mostrar os resultados ao cliente.

Os cientistas da Amazon Search usaram o PyTorch Lightning como uma das principais estruturas para treinar os modelos de aprendizado profundo que potencializam a classificação da pesquisa devido aos recursos de usabilidade adicionados ao PyTorch. O SMDDP não era compatível com modelos de aprendizado profundo escritos no PyTorch Lightning antes do lançamento do novo SageMaker. Isso impediu os cientistas do Amazon Search que preferem usar o PyTorch Lightning de dimensionar seu treinamento de modelo usando técnicas paralelas de dados, diminuindo significativamente o tempo de treinamento e impedindo-os de testar novos experimentos que exigem treinamento mais escalável.

Os primeiros resultados de benchmarking da equipe mostram um tempo de treinamento 7.3 vezes mais rápido para um modelo de amostra quando treinado em oito nós em comparação com uma linha de base de treinamento de nó único. O modelo de linha de base usado nesses benchmarking é uma rede neural perceptron multicamada com sete camadas densas totalmente conectadas e mais de 200 parâmetros. A tabela a seguir resume o resultado do benchmarking em instâncias de treinamento ml.p3.16xlarge do SageMaker.

A seguir, mergulhamos nos detalhes dos novos lançamentos. Se desejar, você pode percorrer nosso correspondente caderno de exemplo.

Execute o PyTorch Lightning com a biblioteca de treinamento distribuído SageMaker

Temos o prazer de anunciar que o SageMaker Data Parallel agora se integra perfeitamente ao PyTorch Lightning no treinamento do SageMaker.

O PyTorch Lightning é uma estrutura de código aberto que fornece uma simplificação para escrever modelos personalizados no PyTorch. De certa forma semelhante ao que Keras fez para o TensorFlow, ou até mesmo Hugging Face, o PyTorch Lightning fornece uma API de alto nível com abstrações para grande parte da funcionalidade de nível inferior do próprio PyTorch. Isso inclui a definição do modelo, criação de perfil, avaliação, poda, paralelismo de modelo, configurações de hiperparâmetro, aprendizado de transferência e muito mais.

Anteriormente, os desenvolvedores do PyTorch Lightning não tinham certeza sobre como migrar perfeitamente seu código de treinamento para clusters de GPU SageMaker de alto desempenho. Além disso, não havia como aproveitar os ganhos de eficiência introduzidos pelo SageMaker Data Parallel.

Para o PyTorch Lightning, em geral, deve haver pouca ou nenhuma alteração de código para simplesmente executar essas APIs no SageMaker Training. Nos cadernos de exemplo, usamos o Estratégia DDP e DDP Plugin métodos.

Há três etapas para usar o PyTorch Lightning com o SageMaker Data Parallel como um back-end otimizado:

1. Use um suporte Contêiner de aprendizado profundo da AWS (DLC) como sua imagem base ou, opcionalmente, crie seu próprio contêiner e instale você mesmo o back-end do SageMaker Data Parallel. Certifique-se de ter o PyTorch Lightning incluído em seus pacotes necessários, como com um requirements.txt arquivo.
2. Faça algumas pequenas alterações de código em seu script de treinamento para habilitar o back-end otimizado. Esses incluem:
   1. Importe a biblioteca SM DDP:

      import smdistributed.dataparallel.torch.torch_smddp
      

   2. Configure o ambiente PyTorch Lightning para o SageMaker:

      from pytorch_lightning.plugins.environments.lightning_environment 
        import LightningEnvironment
      
      env = LightningEnvironment()
      env.world_size = lambda: int(os.environ["WORLD_SIZE"])
      env.global_rank = lambda: int(os.environ["RANK"])

   3. Se você estiver usando uma versão do PyTorch Lightning anterior à 1.5.10, precisará adicionar mais algumas etapas.
      1. Primeiro, adicione a variável de ambiente:

         os.environ["PL_TORCH_DISTRIBUTED_BACKEND"] = "smddp"

      2. Em segundo lugar, certifique-se de usar DDPPlugin, ao invés de DDPStrategy. Se você estiver usando uma versão mais recente, que você pode definir facilmente colocando o requirements.txt no source_dir para o seu trabalho, então isso não é necessário. Veja o seguinte código:

         ddp = DDPPlugin(parallel_devices=[torch.device("cuda", d) for d in range(num_gpus)], cluster_environment=env)

   4. Opcionalmente, defina o back-end do seu grupo de processos como "smddp" no DDPSTrategy objeto. No entanto, se você estiver usando o PyTorch Lightning com o back-end PyTorch DDP, que também é suportado, basta remover este `process_group_backend`parâmetro. Veja o seguinte código:

      ddp = DDPStrategy(
        cluster_environment=env, 
        process_group_backend="smddp", 
        accelerator="gpu")

3. Certifique-se de ter um método de distribuição anotado no estimador, como distribution={"smdistributed":{"dataparallel":{"enabled":True} se você estiver usando o back-end Herring ou distribution={"pytorchddp":{"enabled":True}.

• Para obter uma lista completa de parâmetros adequados no distribution parâmetro, veja nossa documentação SUA PARTICIPAÇÃO FAZ A DIFERENÇA.

Agora você pode iniciar seu trabalho de treinamento do SageMaker! Você pode iniciar seu trabalho de treinamento por meio do Python SDK, Boto3, o console do SageMaker, o Interface de linha de comando da AWS (AWS CLI) e inúmeros outros métodos. Do ponto de vista da AWS, isso é um único comando de API: create-training-job. Se você iniciar este comando de seu terminal local, um AWS Lambda função, um Estúdio Amazon SageMaker notebook, um pipeline KubeFlow ou qualquer outro ambiente de computação depende totalmente de você.

Observe que a integração entre o PyTorch Lightning e o SageMaker Data Parallel atualmente é compatível apenas com as versões mais recentes do PyTorch, a partir da 1.11. Além disso, esta versão está disponível apenas nos DLCs da AWS para SageMaker a partir do PyTorch 1.12. Certifique-se de apontar para esta imagem como sua base. Dentro us-east-1, este endereço é o seguinte:

ecr_image = '763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.12.0-gpu-py38-cu113-ubuntu20.04-sagemaker'

Em seguida, você pode estender seu contêiner do Docker usando isso como sua imagem base ou pode passar isso como uma variável para o image_uri argumento do Estimador de treinamento do SageMaker.

Como resultado, você poderá executar seu código PyTorch Lightning nas GPUs otimizadas do SageMaker Training, com o melhor desempenho disponível na AWS.

Execute o PyTorch Distributed Data Paralelamente no SageMaker

O maior problema Dados Paralelos Distribuídos PyTorch (DDP) resolve é enganosamente simples: velocidade. Uma boa estrutura de treinamento distribuída deve fornecer estabilidade, confiabilidade e, mais importante, excelente desempenho em escala. O PyTorch DDP oferece isso fornecendo aos desenvolvedores de torch APIs para replicar seus modelos em vários dispositivos de GPU, em configurações de nó único e vários nós. A estrutura então gerencia a fragmentação de diferentes objetos do conjunto de dados de treinamento para cada cópia do modelo, calculando a média dos gradientes para cada uma das cópias do modelo para sincronizá-los em cada etapa. Isso produz um modelo na conclusão total da execução completa do treinamento. O diagrama a seguir ilustra esse processo.

O PyTorch DDP é comum em projetos que usam grandes conjuntos de dados. O tamanho preciso de cada conjunto de dados varia muito, mas uma diretriz geral é dimensionar conjuntos de dados, tamanhos de computação e tamanhos de modelo em proporções semelhantes. Também chamado leis de escala, a combinação ideal desses três está em debate e varia de acordo com as aplicações. Na AWS, com base no trabalho com vários clientes, podemos ver claramente os benefícios das estratégias paralelas de dados quando o tamanho geral do conjunto de dados é de pelo menos algumas dezenas de GBs. Quando os conjuntos de dados ficam ainda maiores, implementar algum tipo de estratégia paralela de dados é uma técnica crítica para acelerar o experimento geral e melhorar seu tempo de retorno.

Anteriormente, os clientes que usavam o PyTorch DDP para treinamento distribuído no local ou em outros ambientes de computação não tinham uma estrutura para migrar facilmente seus projetos para o SageMaker Training para aproveitar as GPUs de alto desempenho com um plano de controle contínuo. Especificamente, eles precisavam migrar sua estrutura paralela de dados para o SMDDP ou desenvolver e testar os recursos do PyTorch DDP no SageMaker Training manualmente. Hoje, o SageMaker Training tem o prazer de fornecer uma experiência perfeita para os clientes que integram seu código PyTorch DDP.

Para usar isso de forma eficaz, você não precisa fazer nenhuma alteração em seus scripts de treinamento.

Você pode ver esse novo parâmetro no código a seguir. No distribution parâmetro, basta adicionar pytorchddp e defina ativado como true.

estimator = PyTorch(
    base_job_name="pytorch-dataparallel-mnist",
    source_dir="code",
    entry_point = "my_model.py",
    ... 
    # Training using SMDataParallel Distributed Training Framework
    distribution = {"pytorchddp": {"enabled": "true"}}
)

Essa nova configuração começa nas versões 2.102.0 do SDK do Python do SageMaker e no DLC 1.11 do PyTorch.

Para desenvolvedores de PyTorch DDP que estão familiarizados com o popular tocha framework, é útil saber que isso não é necessário no ambiente de treinamento SageMaker, que já oferece tolerância a falhas robusta. No entanto, para minimizar as reescritas de código, você pode trazer outro script de inicialização que execute esse comando como seu ponto de entrada.

Agora, os desenvolvedores do PyTorch podem mover facilmente seus scripts para o SageMaker, garantindo que seus scripts e contêineres possam ser executados perfeitamente em vários ambientes de computação.

Isso os prepara para, no futuro, aproveitar Bibliotecas de treinamento distribuídas do SageMaker que fornecem topologias de treinamento otimizadas para AWS para oferecer até 40% de aprimoramentos de aceleração. Para desenvolvedores PyTorch, esta é uma única linha de código! Para o código PyTorch DDP, você pode simplesmente definir o back-end para smddp na inicialização (veja Modificar um script de treinamento do PyTorch), conforme mostrado no código a seguir:

import smdistributed.dataparallel.torch.torch_smddp
import torch.distributed as dist
dist.init_process_group(backend='smddp')

Como vimos acima, você também pode definir o backend de DDPStrategy para smddp ao usar Relâmpago. Isso pode levar até 40% de aceleração geral para grandes aglomerados! Para saber mais sobre o treinamento distribuído no SageMaker, consulte nosso webinar sob demandaapoio laptops, relevante documentação e papéis.

Conclusão

Neste post, apresentamos dois novos recursos da família SageMaker Training. Isso torna muito mais fácil para os desenvolvedores do PyTorch usar seu código existente no SageMaker, tanto o PyTorch DDP quanto o PyTorch Lightning.

Também mostramos como o Amazon Search usa o SageMaker Training para treinar seus modelos de aprendizado profundo e, em particular, o PyTorch Lightning com a biblioteca coletiva otimizada do SageMaker Data Parallel como back-end. A mudança para o treinamento distribuído em geral ajudou o Amazon Search a atingir tempos de treinamento 7.3 vezes mais rápidos.

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.

Karan Dhiman é engenheiro de desenvolvimento de software na AWS, com sede em Toronto, Canadá. Ele é muito apaixonado pelo espaço de Machine Learning e pela criação de soluções para acelerar as cargas de trabalho de computação distribuída.

Vishwa Karia é engenheiro de desenvolvimento de software no AWS Deep Engine. Seus interesses estão na interseção de Machine Learning e Sistemas Distribuídos e ela também é apaixonada por capacitar mulheres em tecnologia e IA.

Eiman Elnahrawy é engenheiro de software principal da Amazon Search, liderando os esforços de aceleração, dimensionamento e automação de aprendizado de máquina. Sua experiência abrange várias áreas, incluindo Machine Learning, Sistemas Distribuídos e Personalização.