Repolode
Enviar
google

google/paxml

Pax es un framework basado en JAX para configurar y ejecutar experimentos de machine learning a gran escala, principalmente con modelos de lenguaje, en TPU.

Ver en GitHub
Estrellas
560
Forks
73
Issues abiertos
30
Último commit
hace 12 días

¿Qué es?

Pax (Paxml) es un framework basado en JAX para configurar y ejecutar experimentos de machine learning. Se centra en experimentos totalmente configurables y en la paralelización, y está pensado para entrenar modelos grandes, como modelos de lenguaje Transformer solo con decodificador. Los modelos, las tareas y los learners se definen con hiperparámetros basados en dataclasses, que se apoyan en la biblioteca de capas complementaria Praxis y admiten configuración con Fiddle.

El proyecto está orientado a las TPU VMs de Google Cloud, incluidas las configuraciones v4 de un solo host, los segmentos de TPU Pod y las configuraciones multislice. Incluye configuraciones de ejemplo para entrenar con el conjunto de datos C4 con 1B, 16B y más parámetros, así como una prueba de Model FLOPs Utilization (MFU) en ejecuciones de escalado débil en TPU v4 Pods. Los datos de entrada pueden proceder de SeqIO, Lingvo o subclases personalizadas de BaseInput; también ofrece orientación sobre el sharding multihost y la evaluación.

El repositorio señala que NVIDIA mantiene una versión independiente y actualizada de Pax con soporte para GPU, a la que hace referencia desde su propio repositorio.

Funciones clave

  • 01Experimentos configurables e hiperparámetros basados en dataclasses y Fiddle.
  • 02Compatibilidad con paralelismo mediante pjit/SPMD, pmap y configuraciones multislice de TPU.
  • 03Configuraciones de ejemplo para entrenar con C4 modelos de 1B, 16B y del estilo GPT-3 XL.
  • 04Varias opciones de entrada de datos, incluidos adaptadores de SeqIO y Lingvo, y entradas personalizadas.
  • 05Pruebas de Model FLOPs Utilization en TPU v4 Pods.

Casos de uso

  • Entrenar modelos de lenguaje grandes solo con decodificador en Google Cloud TPU v4 pods.
  • Ejecutar experimentos reproducibles basados en JAX con configuraciones ajustables del modelo y del paralelismo.
  • Comparar la eficiencia del entrenamiento (MFU) con distintos tamaños de modelo y escalas de chips.

Inicio rápido

  1. 1

    Crea una Cloud TPU VM (configura primero las variables de entorno)

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

    Conéctate por SSH a la TPU VM

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

    Instala Pax desde PyPI junto con JAX para TPU

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

    Ejecuta un modelo de prueba de ejemplo

    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 usarlo

  • Está diseñado para hardware Cloud TPU; la configuración y los ejemplos del README presuponen el uso de TPU VMs.
  • La configuración depende de recursos de Google Cloud, como un proyecto de GCP y un bucket de almacenamiento.

Temas

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

Preguntas frecuentes

¿Qué es paxml?
Pax es un framework basado en JAX para configurar y ejecutar experimentos de machine learning a gran escala, principalmente con modelos de lenguaje, en TPU. Pax (Paxml) es un framework basado en JAX para configurar y ejecutar experimentos de machine learning. Se centra en experimentos totalmente configurables y en la paralelización, y está pensado para entrenar modelos grandes, como modelos de lenguaje Transformer solo con decodificador. Los modelos, las tareas y los learners se definen con hiperparámetros basados en dataclasses, que se apoyan en la biblioteca de capas complementaria Praxis y admiten configuración con Fiddle.
¿paxml es gratuito? ¿Qué licencia tiene?
paxml es open source bajo la licencia Apache-2.0. Revisa los términos de la licencia antes de usarlo en productos comerciales.
¿paxml sigue mantenido?
Mantenimiento activo. El último commit de paxml fue el 26 oct 2026.
¿Para qué puedo usar paxml?
Entrenar modelos de lenguaje grandes solo con decodificador en Google Cloud TPU v4 pods. Ejecutar experimentos reproducibles basados en JAX con configuraciones ajustables del modelo y del paralelismo. Comparar la eficiencia del entrenamiento (MFU) con distintos tamaños de modelo y escalas de chips.
¿Cómo se instala paxml?
Crea una Cloud TPU VM (configura primero las variables de entorno): gcloud compute tpus tpu-vm create $TPU_NAME --zone=$ZONE --version=$VERSION --project=$PROJECT --accelerator-type=$ACCELERATOR Conéctate por SSH a la TPU VM: gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE Instala Pax desde PyPI junto con JAX para TPU: python3 -m pip install paxml jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html Ejecuta un modelo de prueba de ejemplo: python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=gs://<your-bucket>
¿Dónde consigo paxml?
El código fuente y las instrucciones de instalación están en GitHub: https://github.com/google/paxml

Resumen generado por IA a partir del README del proyecto. Las estadísticas de GitHub se actualizan a diario.

Proyectos relacionados