Repolode
Enviar
google

google/paxml

Pax é um framework baseado em JAX para configurar e executar experimentos de machine learning em larga escala, principalmente com modelos de linguagem, em TPUs.

Ver no GitHub
Estrelas
560
Forks
73
Issues abertas
30
Último commit
há 12 dias

O que é?

Pax (Paxml) é um framework para configurar e executar experimentos de machine learning com JAX. Ele foi desenvolvido em torno de experimentos totalmente configuráveis e paralelização, com foco no treinamento de modelos grandes, como modelos de linguagem Transformer apenas com decoder. Modelos, tarefas e learners são definidos com hiperparâmetros baseados em dataclasses, que usam a biblioteca de camadas complementar Praxis e oferecem suporte à configuração com Fiddle.

O projeto é voltado para TPU VMs do Google Cloud, incluindo configurações v4 de host único, slices de TPU Pod e configurações multislice. Ele inclui configurações de exemplo para treinar com o dataset C4 em modelos com 1B, 16B e mais parâmetros, além de um benchmark de Model FLOPs Utilization (MFU) em execuções com weak scaling nos TPU v4 Pods. Os dados podem vir de SeqIO, Lingvo ou subclasses personalizadas de BaseInput; o projeto também traz orientações sobre sharding em ambientes multihost e avaliação.

O repositório informa que a NVIDIA mantém uma versão separada e atualizada do Pax com suporte a GPU, indicada em seu próprio repositório.

Recursos principais

  • 01Experimentos configuráveis e hiperparâmetros baseados em dataclasses e Fiddle.
  • 02Suporte a paralelismo com pjit/SPMD, pmap e configurações multislice de TPU.
  • 03Configurações de exemplo para treinamento com C4 em modelos de 1B, 16B e no estilo GPT-3 XL.
  • 04Várias opções de entrada de dados, incluindo SeqIO, adaptadores do Lingvo e entradas personalizadas.
  • 05Benchmarks de Model FLOPs Utilization em TPU v4 Pods.

Casos de uso

  • Treinar modelos de linguagem grandes, apenas com decoder, em Google Cloud TPU v4 Pods.
  • Executar experimentos reproduzíveis baseados em JAX, com configurações personalizáveis de modelo e paralelismo.
  • Medir a eficiência do treinamento (MFU) em diferentes escalas de modelos e chips.

Início rápido

  1. 1

    Crie uma Cloud TPU VM (configure primeiro as variáveis de ambiente)

    gcloud compute tpus tpu-vm create $TPU_NAME --zone=$ZONE --version=$VERSION --project=$PROJECT --accelerator-type=$ACCELERATOR
  2. 2

    Acesse a TPU VM por SSH

    gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
  3. 3

    Instale o Pax pelo PyPI junto com o JAX para TPU

    python3 -m pip install paxml jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
  4. 4

    Execute um modelo de teste de exemplo

    python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=gs://<your-bucket>

Antes de usar

  • Foi desenvolvido para hardware Cloud TPU; a configuração e os exemplos do README pressupõem o uso de TPU VMs.
  • A configuração depende de recursos do Google Cloud, como um projeto GCP e um bucket de armazenamento.

Tópicos

#jax#large-language-models#tpu#distributed-training#parallelism#machine-learning#c4#gpt#llm#model-flops

Perguntas frequentes

O que é paxml?
Pax é um framework baseado em JAX para configurar e executar experimentos de machine learning em larga escala, principalmente com modelos de linguagem, em TPUs. Pax (Paxml) é um framework para configurar e executar experimentos de machine learning com JAX. Ele foi desenvolvido em torno de experimentos totalmente configuráveis e paralelização, com foco no treinamento de modelos grandes, como modelos de linguagem Transformer apenas com decoder. Modelos, tarefas e learners são definidos com hiperparâmetros baseados em dataclasses, que usam a biblioteca de camadas complementar Praxis e oferecem suporte à configuração com Fiddle.
paxml é gratuito? Qual é a licença?
paxml é open source sob a licença Apache-2.0. Confira os termos da licença antes de usá-lo em produtos comerciais.
paxml ainda é mantido?
Manutenção ativa. O último commit de paxml foi em 26 de out. de 2026.
Para que posso usar paxml?
Treinar modelos de linguagem grandes, apenas com decoder, em Google Cloud TPU v4 Pods. Executar experimentos reproduzíveis baseados em JAX, com configurações personalizáveis de modelo e paralelismo. Medir a eficiência do treinamento (MFU) em diferentes escalas de modelos e chips.
Como instalar paxml?
Crie uma Cloud TPU VM (configure primeiro as variáveis de ambiente): gcloud compute tpus tpu-vm create $TPU_NAME --zone=$ZONE --version=$VERSION --project=$PROJECT --accelerator-type=$ACCELERATOR Acesse a TPU VM por SSH: gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE Instale o Pax pelo PyPI junto com o JAX para TPU: python3 -m pip install paxml jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html Execute um modelo de teste de exemplo: python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=gs://<your-bucket>
Onde consigo paxml?
O código-fonte e as instruções de instalação estão no GitHub: https://github.com/google/paxml

Resumo gerado por IA a partir do README do projeto. As estatísticas do GitHub são atualizadas diariamente.

Projetos relacionados