PyTorch Monarch é lançado

• PyTorch Monarch é um novo framework projetado para dar suporte a treinamento e inferência distribuídos eficientes de modelos de grande porte
• Expande a estrutura modular existente do PyTorch, oferecendo recursos para particionar e gerenciar automaticamente redes neurais gigantes entre vários dispositivos e nós
• Por meio de uma API capaz de controlar de forma integrada paralelismo de modelo, paralelismo em pipeline e paralelismo de dados, reduz a carga de configurações complexas para os desenvolvedores
• O Monarch demonstra alta eficiência especialmente em cargas de trabalho intensivas em memória, como grandes modelos de linguagem (LLM) e sistemas de recomendação
• Como parte de uma iniciativa para alcançar ao mesmo tempo escalabilidade e otimização de desempenho dentro do ecossistema PyTorch, ele é um componente central da infraestrutura de treinamento distribuído da próxima geração

---

Visão geral do PyTorch Monarch

• O PyTorch Monarch é apresentado como um novo componente do PyTorch para simplificar o treinamento distribuído e a inferência de modelos de grande porte
  • Foi projetado para manter a flexibilidade existente do PyTorch, ao mesmo tempo em que permite posicionar com eficiência modelos com bilhões de parâmetros em várias GPUs e nós
  • Oferece recursos automatizados de particionamento e otimização de comunicação, sem a necessidade de configurar manualmente estratégias complexas de paralelização
• O principal objetivo do Monarch é elevar o nível de abstração do paralelismo de modelo, permitindo que os desenvolvedores se concentrem na arquitetura do modelo
  • Diferentes técnicas de paralelização, como paralelismo de dados, paralelismo em pipeline e paralelismo de tensor, podem ser controladas por uma única interface unificada
  • Com isso, reduz significativamente a complexidade do código e a sobrecarga de comunicação em comparação com frameworks de treinamento distribuído existentes

Principais recursos e características técnicas

• O Monarch posiciona cada camada do modelo no dispositivo ideal por meio de um algoritmo de particionamento automático
  • A estratégia de particionamento é determinada dinamicamente considerando capacidade de memória da GPU, largura de banda de comunicação e carga computacional
  • Essa automação demonstra alta eficiência especialmente em LLMs, modelos baseados em Transformer e sistemas de recomendação de grande escala
• Fornece uma API unificada de paralelização, permitindo que desenvolvedores experimentem diferentes estratégias de paralelização com uma única base de código
  • Por exemplo, é possível executar o mesmo modelo com uma combinação de paralelismo de dados e paralelismo em pipeline, ou migrar para paralelismo de tensor
  • Essa flexibilidade facilita a busca pela otimização conforme o tamanho do modelo e a configuração de hardware
• O Monarch é compatível com os recursos existentes do PyTorch DistributedDataParallel(DDP) e Fully Sharded Data Parallel(FSDP)
  • É possível migrar para o Monarch sem grandes modificações na base de código existente
  • Também se integra com TorchScript e TorchDynamo do PyTorch, oferecendo suporte à otimização de compilação e execução

Desempenho e casos de uso

• Segundo benchmarks iniciais, o Monarch alcançou melhoria de 20% a 30% na eficiência de comunicação e redução de 15% no uso de memória em comparação com o treinamento distribuído existente do PyTorch
  • Especialmente em modelos com dezenas de bilhões de parâmetros, houve grande melhora na velocidade de treinamento e na utilização de GPU
  • Isso foi validado experimentalmente em grandes modelos de linguagem (por exemplo, da família GPT) e em sistemas de recomendação
• O Monarch funciona tanto em ambientes de nuvem quanto on-premises e é compatível com as principais infraestruturas de nuvem, como AWS, Azure e GCP
  • Também oferece integração com frameworks de nível superior, como PyTorch Lightning e Hugging Face Transformers

Significado no ecossistema PyTorch

• O Monarch é visto como uma expansão de infraestrutura central do PyTorch para responder à era dos modelos de IA em grande escala
  • Rompe com o paradigma tradicional de treinamento centrado em uma única GPU e fornece a base para treinamento de modelos gigantes usando milhares de GPUs
  • Atua como uma solução padronizada de treinamento distribuído que permite a pesquisadores e empresas obter ao mesmo tempo escalabilidade e eficiência
• A equipe do PyTorch planeja lançar o Monarch como open source e continuar evoluindo-o com base no feedback da comunidade
  • No futuro, estão previstos recursos como otimização automática, agendamento dinâmico e paralelismo híbrido
  • Como framework distribuído de próxima geração do PyTorch, a expectativa é que contribua para a democratização da infraestrutura de IA e para a melhoria da acessibilidade

[GN⁺]

Comentários do Hacker News

• Parece que este projeto mira uma camada diferente da do Tinker
  Pelo post de apresentação do Tinker, o Tinker é um serviço gerenciado de fine-tuning, enquanto o Monarch fornece primitivos de infraestrutura
  Então fico curioso se daria para construir um serviço como o Tinker em cima do Monarch

  • No momento, algo como o TorchForge já está rodando em cima dele

• Parece que a oxidação do PyTorch começou
  O Monarch é dividido entre um frontend baseado em Python e um backend implementado em Rust
  No geral, parece um projeto bem interessante

  • Segundo várias fontes, o Monarch é um framework experimental do PyTorch, não um substituto
    Então ainda deve dar para curtir grafos cíclicos baseados em std::shared_ptr e vazamentos de memória
    É uma pena que não tenham reescrito tudo do zero em uma linguagem funcional
  • Isso não parece ser uma versão oxidada de um projeto existente, mas sim um projeto totalmente novo

• Eu mesmo já criei uma extensão do PyTorch — mycelya-torch
  Minha versão ainda não suporta comunicação entre nós, mas achei interessante a forma como o Monarch busca desempenho
  O Monarch provavelmente usa cloudpickle para compartilhar o código com todos os nós, o que é eficiente porque só tem custo na configuração inicial
  Também me chamou atenção a estrutura de fan-out de mensagens a partir de um único controlador
  Mas fico curioso sobre o suporte a kernels customizados e o nível de granularidade no controle da comunicação entre atores
  No geral, gosto mais dessa abordagem do que de múltiplos controladores

  • Para usar kernels customizados, talvez seja preciso modificar um pouco o código de inicialização dos workers remotos
    Mas dá para embutir (bake-in) diretamente os kernels ou o código de sistema necessários

• Isso parece ter uma estrutura parecida com a do Ray

  • O exemplo de código é quase idêntico ao do Ray
    O Actor do Monarch e as classes com @ray.remote do Ray seguem o mesmo padrão
  • Fico curioso sobre por que usar Monarch em vez de Ray — talvez por uma integração mais estreita com o PyTorch ou com a abstração de tensores
  • O Dask também faz processamento distribuído semelhante, mas foi pensado originalmente para HPC, então o suporte a GPU é limitado
    Veja o site oficial do Dask
  • Pensei a mesma coisa. Principalmente vendo o recente anúncio de colaboração entre PyTorch e Ray, isso fica ainda mais forte
    Blog relacionado

• Interessante, isso essencialmente me lembra o conceito de coarray do Fortran (2008)

  • Ou talvez Hadoop (2006)
    Só que acho bem melhor por não precisar usar MapReduce nem Fortran diretamente

• Havia uma frase dizendo “evita gargalos de host único e usa a malha inteira como um cluster distribuído para passagem de mensagens”,
  se alguém com acesso para editar vir isso, eu gostaria que adicionasse alguns números de referência

• Parece um bom projeto
  Tenho algumas dúvidas

  • Isso é parecido com openMPI?
  • Como essa malha é formada? Só funciona dentro do mesmo host?

• Isso pode virar um projeto importante no mundo de coarrays, mas já há sinais de problema
  O engine de tensores está preso a CUDA e RDMA (ibverbs), e como o código usa GPUDirect RDMA,
  no fim a dependência de CUDA deve ficar ainda mais forte
  Teria sido melhor se tivessem usado OpenUCX

• Parece funcionalmente mais fraco que o Jax
  O Jax usa um compilador poderoso para otimizar a comunicação entre nós

  • Mas o Monarch foca no paradigma de controlador único, enquanto o Jax segue uma estrutura SPMD com múltiplos controladores
    Um controlador único é mais fácil de entender, e múltiplos controladores são mais adequados para certos fluxos de dados
    Também há tentativas interessantes de misturar as duas abordagens