Imágenes de SO y Docker

Google Cloud proporciona imágenes que contienen sistemas operativos, frameworks, bibliotecas y controladores comunes. Google Cloud optimiza estas imágenes preconfiguradas para admitir tus cargas de trabajo de Inteligencia Artificial (IA) y aprendizaje automático (AA).

En este documento, se proporciona una descripción general de las imágenes que usas para implementar, administrar y ejecutar cargas de trabajo en tu entorno de AI Hypercomputer.

Comprende las categorías de imágenes

Las imágenes se agrupan en las siguientes categorías:

  • Frameworks y bibliotecas de IA y AA: Imágenes de Docker preconfiguradas con archivos binarios para frameworks y bibliotecas de AA que simplifican la creación, el entrenamiento y el uso de modelos de AA En AI Hypercomputer, puedes usar imágenes de Docker de Deep Learning Software Layer (DLSL) para ejecutar modelos de AA, como NeMO y MaxText, en un clúster de Google Kubernetes Engine (GKE).
  • Implementación y organización de clústeres: Imágenes del sistema operativo (SO) que usas para implementar y administrar la infraestructura optimizada para el rendimiento en la que se ejecutan tus cargas de trabajo de IA. La imagen que debes usar depende de la infraestructura de GPU que utilices (GPU agrupadas o GPU generales) y de cómo planees implementar tus clústeres:

Frameworks y bibliotecas de IA y AA

Google Cloud proporciona imágenes de Docker que empaquetan frameworks y bibliotecas populares de IA y AA. Estas imágenes proporcionan el software necesario para simplificar el desarrollo, el entrenamiento y la implementación de modelos en tus clústeres optimizados para IA que se ejecutan en AI Hypercomputer.

Imágenes generadas por IA de JAX

Las imágenes de IA de JAX (JAII, anteriormente conocidas como imágenes de pila estable de JAX) para Google Cloud TPU y GPU ofrecen imágenes de Docker listas para usar que contienen el framework de JAX, una colección seleccionada de bibliotecas compatibles y la configuración para la infraestructura de Google Cloud . Las imágenes de TPU de IA de JAX vienen preconfiguradas con bibliotecas de JAX y bibliotecas de TPU. Las imágenes de GPU de IA de JAX vienen preconfiguradas con bibliotecas de JAX y bibliotecas de CUDA/NVIDIA pertinentes.

Diagrama de la imagen generada por IA de JAX

La capa de hardware

Las imágenes de IA de JAX se encuentran en la parte superior de la capa de hardware, que consta de los aceleradores (TPU o GPU) y sus VMs asociadas. Para usar una imagen de IA de JAX, debes aprovisionar VMs de TPU o GPU. Puedes hacerlo con la API de TPU, la API de Compute Engine o la API de GKE.

La capa de framework

La capa de framework proporciona herramientas y bibliotecas para compilar cargas de trabajo de AA. Las imágenes de IA de JAX proporcionan una base preconfigurada para las cargas de trabajo de AA basadas en JAX, incluida la biblioteca principal de JAX y otras dependencias esenciales, para garantizar una experiencia de desarrollo coherente y de alto rendimiento.

La capa LibTPU en la imagen de IA de JAX se compila y se incluye específicamente con la versión de JAX correspondiente. Si usas una versión diferente de JAX, es posible que se produzcan errores o comportamientos inesperados.

La capa de CUDA en la imagen de IA de JAX incluye componentes administrados por NVIDIA, como la imagen de aprendizaje profundo de CUDA de NGC, que se usa como imagen base de la imagen de entrenamiento de GPU. La imagen de GPU también contiene Transformer Engine, una biblioteca personalizada de NVIDIA para acelerar los modelos de Transformer en las GPUs de NVIDIA.

Es posible que se requieran paquetes adicionales específicos de la aplicación, más allá de los que se proporcionan en la imagen de IA de JAX, para tu carga de trabajo de aprendizaje automático específica.

Bibliotecas en JAX AI Images:

Imágenes de TPU

Funcionalidad Nombre del paquete
Bibliotecas o componentes principales
Framework de AA JAX y JAX lib
LibTPU Versión estable de Cloud LibTPU
Biblioteca de capas y modelos Flax
Creación de puntos de control Orbax
Cuantización Qwix
Optimizadores Optax

aqtp

Kernels personalizados Tokamax
Configuración Fiddle
Canalización de entrada tf.data

PyGrain array-record

Generación de perfiles y depuración TensorBoard
Utils Utilidades de bucle comunes

Medición del buen rendimiento del AA Colecciones de AA

Herramientas específicas de la nube
Google Cloud

Google Cloud SDK Google Cloud storage

Diagnóstico

Diagnóstico del acelerador de Cloud Diagnóstico de Cloud TPU

Imágenes de GPU

Funcionalidad Nombre del paquete
Bibliotecas o componentes principales
Framework de AA JAX y JAX lib
Bibliotecas de NVIDIA CUDA Imagen de DL de CUDA
Biblioteca de capas y modelos Flax
Creación de puntos de control Orbax
Optimizadores

Optax aqtp TransformerEngine

Configuración Fiddle
Canalización de entrada

tf.data PyGrain array-record

Generación de perfiles y depuración TensorBoard
Utils

Utilidades comunes de bucle Medición del buen rendimiento del AA Colecciones de AA

Herramientas específicas de la nube
Google Cloud

Google Cloud SDK Google Cloud storage

Diagnóstico Diagnóstico de Cloud Accelerator

Imágenes generadas por IA de JAX actuales

Imágenes de TPU

Imagen generada por IA de JAX Fecha de lanzamiento
JAX 0.9.0, revisión 1 2026-02-03
Revisión 1 de JAX 0.8.2 2026-01-14
JAX 0.8.1, revisión 1 2025-11-21
JAX 0.8.0, revisión 1 2025-10-28
JAX 0.7.2, revisión 1 2025-09-30
JAX 0.7.0, revisión 1 2025-07-29
JAX 0.6.1, revisión 1 2025-06-05
JAX 0.5.2, revisión 2 2025-04-25
JAX 0.5.2, revisión 1 2025-03-17
JAX 0.4.37, revisión 1 2024-12-12
JAX 0.4.35, revisión 1 2024-10-30

Imágenes de GPU

Imagen generada por IA de JAX Fecha de lanzamiento
JAX 0.7.2 con CUDA DL 25.06, revisión 1 2025-09-30
JAX 0.6.1 con CUDA DL 25.03, revisión 1 2025-06-05
JAX 0.5.1 con CUDA DL 25.02, revisión 1 2025-03-17

La capa de aplicación

Implementas tus cargas de trabajo de AA específicas en la capa de la aplicación, que se encuentra sobre la capa del framework. La capa de aplicación contiene el código, los modelos y la lógica específicos de tu aplicación, todos compilados con las herramientas y las bibliotecas que proporciona la capa de framework.

Si bien esta imagen proporciona una base sólida y bien probada para las cargas de trabajo de IA basadas en JAX, es posible que debas agregar dependencias específicas de la aplicación. Cuando lo hagas, te recomendamos que lo hagas de una manera que minimice la interferencia con la capa base preconfigurada, que incluye JAX y sus dependencias principales. La introducción de dependencias a nivel de la aplicación que anulan o entran en conflicto con las dependencias existentes puede causar efectos secundarios, como los siguientes:

  • Comportamientos inesperados: Es posible que tus cargas de trabajo de AA muestren un comportamiento diferente al que tenían antes de que introdujeras dependencias adicionales en la imagen de AA de JAX.
  • Regresiones de rendimiento: Anular las bibliotecas optimizadas relacionadas con JAX puede afectar negativamente los beneficios de rendimiento que proporciona la imagen generada por IA de JAX.
  • Problemas de estabilidad: Los conflictos entre las dependencias que agregaste y las dependencias principales de JAX pueden generar inestabilidad y errores de tiempo de ejecución en tu aplicación.

Frecuencia de actualización

Inicialmente, las imágenes de IA de JAX se proporcionan trimestralmente con el objetivo a corto plazo de tener un programa de lanzamiento sincronizado con cada lanzamiento de JAX. Esto garantiza que puedas aprovechar las funciones y mejoras más recientes tan pronto como estén disponibles.

Asistencia

Cada lanzamiento de imágenes generadas por IA de JAX se adhiere a un ciclo de vida de asistencia técnica por tiempo limitado. En este plazo, abordamos categorías específicas de solicitudes de modificaciones a imágenes existentes generadas por IA de JAX:

  • Vulnerabilidades de seguridad: Priorizamos abordar las vulnerabilidades de seguridad que se descubren en las imágenes base o las dependencias de las imágenes de Docker de JAX Stable Stack. Se lanzarán imágenes actualizadas para mitigar los riesgos potenciales.
  • Cambios rotundos: En caso de que haya cambios rotundos significativos en las bibliotecas o los frameworks subyacentes que usa la imagen de IA de JAX, Google Cloudevalúa e implementa las actualizaciones necesarias para mantener la compatibilidad. Esto puede implicar volver a compilar imágenes de Docker con dependencias actualizadas.

Cuando se descubre una vulnerabilidad o un error de seguridad en una biblioteca dentro de la JAII, incorporamos la biblioteca actualizada en la JAII y fijamos todas las demás versiones de la biblioteca para mantener la estabilidad general. Esto da como resultado una nueva revisión de la JAII.

Cambio mínimo para las revisiones:

Si se encuentra un error en el paquete "X" dentro de JAX-0.4.30-rev1, actualizaremos "X" a su próxima versión (por ejemplo, v2.0) y trataremos de mantener todos los demás paquetes sin cambios. Esto generará una nueva revisión: JAX-0.4.30-rev2, que se lanzará lo antes posible.

Imágenes de Docker con capa de software de aprendizaje profundo (DLSL)

Estas imágenes empaquetan NVIDIA CUDA, NCCL, un framework de AA y un modelo. Proporcionan un entorno listo para usar para cargas de trabajo de aprendizaje profundo. Estas imágenes de Docker de DLSL prediseñadas funcionan a la perfección con tus clústeres de GKE porque las probamos y verificamos durante las pruebas internas de reproducibilidad y regresión.

Las imágenes de Docker de DLSL proporcionan los siguientes beneficios:

  • Software preconfigurado: Las imágenes de Docker de DLSL replican la configuración que se usa para las pruebas internas de reproducibilidad y regresión. Estas imágenes proporcionan un entorno preconfigurado, probado y optimizado, lo que ahorra una gran cantidad de tiempo y esfuerzo en la instalación y la configuración.
  • Administración de versiones: Las imágenes de Docker de DLSL se actualizan con frecuencia. Estas actualizaciones de versión proporcionan la versión estable más reciente de los frameworks y los controladores, y también abordan los parches de seguridad.
  • Compatibilidad con la infraestructura: Las imágenes de Docker de DLSL se compilan y prueban para que funcionen sin problemas con los tipos de máquinas con GPU disponibles en AI Hypercomputer.
  • Instrucciones de inicio rápido: Algunas imágenes de Docker de DLSL incluyen recetas de muestra que te muestran cómo iniciar tus cargas de trabajo que usan las imágenes preconfiguradas.

Complemento de gIB de NeMo + PyTorch + NCCL

Estas imágenes de Docker se basan en la imagen de NVIDIA NeMo NGC. Contienen el complemento gIB de NCCL de Google y agrupan todos los archivos binarios de NCCL necesarios para ejecutar cargas de trabajo en cada máquina aceleradora compatible. Estas imágenes también incluyen herramientas de Google Cloudcomo gcsfusey la gcloud CLIpara implementar cargas de trabajo en Google Kubernetes Engine.

Versión con imágenes de DLSL Versión de las dependencias Series de máquinas Fecha de lanzamiento Fecha de finalización de la asistencia Nombre de la imagen de DLSL
nemo25.04-gib1.0.6-A4
  • NeMo NGC:25.04.01
  • NCCL giB plugin: 1.0.6
A4 3 de julio de 2025 3 de julio de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo25.04-gib1.0.6-A4
nemo25.04-gib1.0.6-A3U
  • NeMo NGC:25.04.01
  • NCCL giB plugin: 1.0.6
A3 Ultra 3 de julio de 2025 3 de julio de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo25.04-gib1.0.6-A3U
nemo25.02-gib1.0.5-A4
  • NeMo NGC:25.02
  • NCCL giB plugin: 1.0.5
A4 14 de marzo de 2025 14 de marzo de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo25.02-gib1.0.5-A4
nemo24.07-gib1.0.2-A3U
  • NeMo NGC:24.07
  • NCCL giB plugin: 1.0.2
A3 Ultra 2 de febrero de 2025 2 de febrero de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo24.07-gib1.0.2-A3U
nemo24.07-gib1.0.3-A3U
  • NeMo NGC:24.07
  • NCCL giB plugin: 1.0.3
A3 Ultra 2 de febrero de 2025 2 de febrero de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo24.07-gib1.0.3-A3U
nemo24.12-gib1.0.3-A3U
  • NeMo NGC:24.12
  • NCCL giB plugin: 1.0.3
A3 Ultra 7 de febrero de 2025 7 de febrero de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo24.12-gib1.0.3-A3U
nemo24.07-tcpx1.0.5-A3Mega
  • NeMo NGC:24.07
  • GPUDirect-TCPX: 1.0.5
A3 Mega 12 de marzo de 2025 12 de marzo de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo24.07-tcpx1.0.5-A3Mega
nemo24.07-tcpx1.0.5-A3High
  • NeMo NGC:24.07
  • GPUDirect-TCPX: 1.0.5
A3 High (8 GPUs) 12 de marzo de 2025 12 de marzo de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo24.07-tcpx1.0.5-A3High

NeMo y PyTorch

Esta imagen de Docker se basa en la imagen de NVIDIA NeMo NGC y, además, incluye Google Cloud herramientas como gcsfuse y gcloud CLI para implementar cargas de trabajo en Google Kubernetes Engine.

Versión con imágenes de DLSL Versión de las dependencias Series de máquinas Fecha de lanzamiento Fecha de finalización de la asistencia Nombre de la imagen de DLSL
nemo24.07--A3U NeMo NGC:24.07 A3 Ultra December 19, 2024 19 de diciembre de 2025 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo:nemo24.07-A3U
nemo24.07-tcpx1.0.5-A3Mega
  • NeMo NGC:24.07
  • GPUDirect-TCPX: 1.0.5
A3 Mega 12 de marzo de 2025 12 de marzo de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo:nemo24.07-tcpx1.0.5-A3Mega
nemo24.07-tcpx1.0.5-A3High
  • NeMo NGC:24.07
  • GPUDirect-TCPX: 1.0.5
A3 High (8 GPUs) 12 de marzo de 2025 12 de marzo de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo:nemo24.07-tcpx1.0.5-A3High

MaxText + caja de herramientas de JAX

Esta imagen de Docker se basa en la imagen de la caja de herramientas de NVIDIA JAX y contieneGoogle Cloud herramientas como gcsfuse y gcloud CLI para implementar cargas de trabajo en Google Kubernetes Engine.

Versión con imágenes de DLSL Versión de las dependencias Series de máquinas Fecha de lanzamiento Fecha de finalización de la asistencia Nombre de la imagen de DLSL
toolbox-maxtext-2025-01-10-A3U JAX toolbox: maxtext-2025-01-10 A3 Ultra 11 de marzo de 2025 11 de marzo de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-maxtext-gpu:toolbox-maxtext-2025-01-10-A3U
jax0.5.1-cu12-A3Mega
  • JAX: 0.5.1
  • CUDA: 12.x
  • MaxText: Latest
A3 Mega 17 de marzo de 2025 17 de marzo de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-gpu-maxtext:jax0.5.1-cu12-A3Mega
jax0.5.1-cu12-A3High
  • JAX: 0.5.1
  • CUDA: 12.x
  • MaxText: Latest
A3 High (8 GPUs) 17 de marzo de 2025 17 de marzo de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-gpu-maxtext:jax0.5.1-cu12-A3High

Pila estable de MaxText y JAX

Esta imagen de Docker se basa en la pila estable de JAX y MaxText. Esta imagen también incluye dependencias, como dnsutils, para ejecutar cargas de trabajo en Google Kubernetes Engine.

Versión con imágenes de DLSL Versión de las dependencias Series de máquinas Fecha de lanzamiento Fecha de finalización de la asistencia Nombre de la imagen de DLSL
jax-maxtext-gpu:jax0.5.1-cuda_dl25.02-rev1-maxtext-20150317
  • JAX Stable stacks:jax0.5.1-cuda_dl25.02-rev1
  • Confirmación de maxtext: 54e98c9e62caa426cf5902be068533ddb4fb79f5
A4 17 de marzo de 2025 17 de marzo de 2026 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-maxtext-gpu:jax0.5.1-cuda_dl25.02-rev1-maxtext-20150317

Implementación y organización de clústeres

Las imágenes de SO incluyen todos los componentes de software necesarios para implementar un sistema operativo en una instancia de procesamiento o un nodo de GKE. El sistema operativo administra los recursos de hardware subyacentes, como los aceleradores y las redes. Esto proporciona los recursos de procesamiento para tu carga de trabajo de IA.

Imágenes de nodos de GKE

GKE implementa clústeres con imágenes de nodos. Estas imágenes de nodos están disponibles para varios sistemas operativos, como Container-Optimized OS, Ubuntu y Windows Server. Las imágenes de nodo de Container-Optimized OS con containerd (cos_containerd) que necesitas para implementar clústeres de GKE Autopilot incluyen optimizaciones para admitir tus cargas de trabajo de IA y AA.

Para obtener más información sobre estas imágenes de nodo, consulta Imágenes de nodo.

Imágenes de SO de Slurm

Los clústeres de Slurm implementan nodos de procesamiento y de controlador como instancias de procesamiento en Compute Engine.

Para aprovisionar clústeres de Slurm optimizados para IA, debes usar Cluster Toolkit. De forma predeterminada, los planos del clúster para los clústeres A4X, A4 y A3 Ultra implementan instancias de procesamiento que usan imágenes de procesamiento avanzadas prediseñadas. Para otras series de máquinas o cuando se compilan imágenes personalizadas de Slurm, el esquema del clúster compila automáticamente una imagen de SO personalizada durante la implementación que instala el software de sistema requerido para la administración del clúster y la carga de trabajo en los nodos de Slurm. Puedes modificar los planos predeterminados antes de implementarlos para personalizar el software que incluyen tus imágenes.

En la siguiente sección, se resume el software que el esquema del clúster instala en tus nodos de Slurm A4, A3 Ultra, A3 Mega y A3 High (8 GPUs). Las plantillas de clúster extienden las imágenes de SO del acelerador con Ubuntu LTS.

A4X Max

El plan de A4X Max disponible en GitHub incluye el siguiente software de forma predeterminada: