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.

Aqui está uma lista de palestras do Ray Summit 2023.

Alguns dos meus favoritos são

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á!


Artigo Original