A stack de IA moderna - Ray
Uma estrutura distribuída de código aberto usada por Uber, Spotify, Instacart, Netflix, Cruise, ByteDance e OpenAI
24 de Setembro, 2024
História
Em 2016, alunos e pesquisadores trabalharam em um projeto de classe no Berkeley Rise Lab para dimensionar o treinamento de redes neurais distribuídas e o aprendizado por reforço.
Isso deu origem ao artigo Ray: A Distributed Framework for Emerging AI Applications.
Hoje, o Ray é uma estrutura de computação distribuída de código aberto para produzir e dimensionar cargas de trabalho de ML Python de forma simples.
Ele resolve três desafios principais no ML distribuído.
- Remover restrição de computação: acesso remoto à computação virtualmente infinita
- Tolerante a falhas: redireciona automaticamente as tarefas com falha para outras máquinas no cluster
- Gerenciamento de estado: compartilhe dados entre tarefas e coordene entre dados
Mas antes de nos aprofundarmos, vamos entender por que precisamos de Ray.
Crise de hardware
Com o LLM e o GenAI Boom, há uma lacuna crescente entre a demanda e a oferta de computação.
Vejamos este gráfico abaixo.
Ray, uma estrutura distribuída unificada para a pilha de IA moderna
Vemos que as demandas de computação para treinamento de sistemas de ML têm crescido 10 vezes a cada 18 meses.
Há uma enorme lacuna entre o treinamento de modelos SOTA e o desempenho de um único núcleo. Embora o hardware especializado forneça um aumento impressionante no desempenho, ele ainda fica aquém de satisfazer as demandas de computação, e essa lacuna só crescerá exponencialmente.
Mesmo que os tamanhos dos modelos parem de crescer, levará décadas para que o hardware especializado se recupere.
A melhor solução agora é distribuir cargas de trabalho de IA.
Mas isso vem com seus desafios.
Desafios de um aplicativo de IA
A criação de aplicativos de IA exige que os desenvolvedores unam cargas de trabalho de ingestão de dados, pré-processamento, treinamento, ajuste fino, previsão e serviço.
Isso é desafiador, pois cada carga de trabalho requer sistemas diferentes, cada um com suas próprias APIs, semânticas e restrições.
Ray, uma estrutura distribuída unificada para a pilha de IA moderna
Com o Ray, você tem um sistema para suportar todas essas cargas de trabalho
Pilha de bibliotecas Ray — kit de ferramentas unificado para cargas de trabalho de ML.
De seus documentos:
Cada uma das cinco bibliotecas nativas de Ray distribui uma tarefa específica de ML:
- Dados: carregamento e transformação de dados escalonáveis e independentes de estrutura em treinamento, ajuste e previsão.
- Treinar: treinamento de modelo distribuído de vários nós e vários núcleos com tolerância a falhas que se integra a bibliotecas de treinamento populares.
- Ajustar: ajuste de hiperparâmetro escalável para otimizar o desempenho do modelo.
- Servir: Servidor escalável e programável para implantar modelos para inferência on-line, com microbatching opcional para melhorar o desempenho.
- RLlib: cargas de trabalho de aprendizado por reforço distribuídas escalonáveis.
Como as empresas estão usando o Ray
Ray, uma estrutura distribuída unificada para a pilha de IA moderna
A OpenAI usou o Ray para coordenar o treinamento do ChatGPT.
A Cohere usou o Ray para treinar seus LLMs em escala junto com PyTorch, JAX e TPU.
Abaixo está uma imagem da Alpa, usando o Ray para agendar GPUs para treinamento distribuído.
Como a Ray resolve os desafios de infraestrutura de IA generativa e LLM
Ray resolve os dois desafios mais comuns para o treinamento distribuído de modelos generativos:
- Como particionar efetivamente o modelo em vários aceleradores?
- Como configurar seu treinamento para ser tolerante a falhas em instâncias preemptivas?
Empresas como Shopify, Spotify, Pinterest e Roblox usam o Ray para dimensionar sua infraestrutura de ML.
A Shopify usa o Ray em sua plataforma Merlin para simplificar os fluxos de trabalho de ML, desde a prototipagem até a produção, utilizando o Ray Train and Tune para treinamento distribuído e ajuste de hiperparâmetros.
O Spotify emprega o Ray para treinamento e ajuste de modelos paralelos para otimizar seus sistemas de recomendação, enquanto o Pinterest utiliza o Ray para processamento eficiente de dados e gerenciamento de infraestrutura escalável.
Na Roblox, Ray facilita a inferência de IA em larga escala em ambientes de nuvem híbrida, permitindo soluções de ML robustas e escaláveis
O núcleo de Ray
API minimalista
A API principal do Ray é de apenas 6 chamadas.
ray.init()
@ray.remote
def big_function():
...
futures = slow_function.remote() # invoke
ray.get(futures) # return an object
ray.put() # store object in object store
ray.wait() # get objects that are ready
ray.shutdown()
Aqui está um exemplo da classe Python Counter se tornando uma função assíncrona.
Ray, uma estrutura distribuída unificada para a pilha de IA moderna
Dando uma chance a Ray
Depois de ver o que Ray pode fazer, decidi tentar.
Digamos que temos essa função que calcula a soma de todos os números primos até n.is_prime
def is_prime(n):
if n < 2:
return False
for i in range(2, int(math.sqrt(n)) + 1):
if n % i == 0:
return False
return True
def sum_primes(limit):
return sum(num for num in range(2, limit) if is_prime(num))
Vamos testar o Python normal para a soma de até 10 milhões e realizar esse cálculo 8 vezes.
%%time
# Sequential execution
n_calculations = 8
limit = 10_000_000
sequential_results = [sum_primes(limit) for _ in range(n_calculations)]
# CPU times: user 12min 31s, sys: 7.56 s, total: 12min 39s
# Wall time: 12min 58s
Isso levou 13 minutos no total!
Vamos ver o quanto Ray pode acelerar isso.
%%time
# Parallel execution
futures = [sum_primes.remote(limit) for _ in range(n_calculations)]
parallel_results = ray.get(futures)
# CPU times: user 477 ms, sys: 366 ms, total: 843 ms
# Wall time: 4min 2s
É uma melhoria de 3x!
Se você clicar no painel, poderá ver o andamento dos trabalhos.
Como fica no painel
Estou apenas começando com Ray.
Fique ligado no meu próximo artigo sobre incorporações de computação com Ray Data!
Quer saber mais sobre Ray?
Assista a esta palestra de 20 minutos Ray, uma estrutura distribuída unificada para a pilha de IA moderna por Ion Stoica, cofundador da Anyscale e DataBricks.
Mais fundo na toca 🕳️🐇 do coelho
Se você quiser um mergulho profundo na arquitetura do Ray, confira o whitepaper
O blog Anyscale tem uma tonelada de casos de uso e tutoriais interessantes.
- Ajuste fino de Llama-3, Mistral e Mixtral com qualquer escala
- Construindo aplicativos LLM baseados em RAG para produção
- RAG em escala: computação de incorporação 10x mais barata com Anyscale e Pinecone
- Lições dos modelos Stable Diffusion v2 pré-treinamento em 2 bilhões de imagens
Aqui está uma lista de palestras do Ray Summit 2023.
Alguns dos meus favoritos são
- Cluster de treinamento heterogêneo com Ray na Netflix
- Gráficos de inferência no LinkedIn usando Ray-Serve
- Construa a plataforma de treinamento Instacart no Ray
- Construindo a plataforma de aprendizado de máquina do Samsara com Ray
- Ray Train: uma biblioteca pronta para produção para aprendizado profundo distribuído
Falando da cúpula,
Junte-se ao Ray Summit 2024!
https://raysummit.anyscale.com
O Ray Summit 2024 acontecerá na próxima semana, de 30 de setembro a 2 de outubro, em São Francisco!
Se você está em SF e pensando em participar, compre um ingresso com o código promocional RodolfoY15 para um ótimo desconto!
Eu vou comparecer, então, por favor, diga oi se você me pegar lá!