Entrenamiento de Qwen3 en AWS Trainium con torchtitan y kernels personalizados

AWS ha publicado una detallada guía técnica que demuestra cómo entrenar el modelo de lenguaje grande Qwen3 en sus aceleradores Trainium2 utilizando el framework torchtitan. Publicado en octubre de 2026, el documento describe tres rutas de configuración distintas para ingenieros que buscan aprovechar las interfaces nativas de PyTorch en hardware especializado. Esta documentación sirve como referencia práctica para equipos de machine learning que desean optimizar los flujos de trabajo de pre-entrenamiento y ajuste fino (fine-tuning) en la infraestructura de Amazon Web Services.

Qué ocurrió

El artículo ofrece una guía de implementación paso a paso para ejecutar modelos Qwen3 en una instancia trn2.48xlarge, que contiene 16 aceleradores Trainium2 y 128 NeuronCores físicos. Los autores demuestran que los desarrolladores pueden usar torchtitan en Trainium de manera similar a como lo harían en clústeres GPU estándar, gracias a que TorchNeuron registra Trainium como un dispositivo nativo de PyTorch. Esto permite al framework detectar el hardware al iniciar y compilar operaciones ya sea de forma inmediata (eagerly) o mediante torch.compile.

La guía cubre tres configuraciones específicas de entrenamiento. El primer método utiliza Hugging Face Transformers v5 como backend de modelado, permitiendo a los usuarios cargar modelos directamente desde el Hub usando AutoModelForCausalLM. El segundo enfoque cambia a la implementación nativa de Qwen3 de torchtitan, ofreciendo mayor control sobre la arquitectura del modelo y las estrategias de fragmentación (sharding). El tercer método introduce personalización avanzada mediante la integración de kernels mixture-of-experts de la Neuron Kernel Interface (NKI), lo que permite a los desarrolladores adjuntar operadores personalizados para capas críticas en cuanto a rendimiento.

Cada configuración utiliza el mismo patrón de lanzamiento torchrun, asegurando consistencia entre diferentes entornos experimentales. La guía enfatiza la reproducibilidad especificando versiones exactas de software, incluyendo PyTorch 2.14.0, Transformers 5.16.0 y torchtitan-neuron release-2.33. También detalla cómo gestionar los checkpoints, verificando que los pesos guardados en formato safetensors de Hugging Face puedan cargarse nuevamente sin claves faltantes, confirmando así la integridad del pipeline de entrenamiento.

Cómo funciona

El mecanismo central se basa en la integración entre torchtitan y la pila AWS Neuron. Cuando se instala torch_neuronx, este expone Trainium como un dispositivo neuron. Torchtitan luego utiliza primitivas estándar de entrenamiento distribuido de PyTorch, como Fully Sharded Data Parallel (FSDP) y paralelismo de tensores, para distribuir la carga de trabajo entre los 64 NeuronCores lógicos disponibles en la instancia. El framework gestiona la fragmentación de parámetros del modelo, gradientes y estados del optimizador, mientras que el paralelismo de secuencia maneja la fragmentación de activaciones a lo largo de la dimensión de secuencia.

Para la integración de kernels personalizados, la guía demuestra cómo utilizar el registro de anulación (override registry) de torchtitan. Los desarrolladores pueden escribir kernels en NKI, el lenguaje de kernel de Neuron, y envolverlos como operadores de PyTorch utilizando wrap_nki y nki_op. Estos operadores se registran luego para reemplazar módulos específicos dentro del grafo del modelo. Por ejemplo, un kernel personalizado de embeddings rotacionales puede reemplazar la implementación predeterminada de PyTorch. El sistema valida estas anulaciones al inicio, asegurando que cada reemplazo apunte a los nombres completamente cualificados correctos dentro de la estructura del modelo.

El bucle de entrenamiento informa métricas clave como la pérdida (loss), la norma del gradiente, el uso de memoria de alto ancho de banda y la utilización de FLOPs del modelo. Esta visibilidad permite a los ingenieros monitorear el consumo de recursos y ajustar dinámicamente la configuración de paralelismo. El proceso de compilación ocurre durante el primer paso de entrenamiento, el cual tarda más que los pasos subsiguientes, pero habilita una ejecución optimizada para el resto de la ejecución.

Detalles clave

• La configuración requiere una instancia trn2.48xlarge con 64 NeuronCores lógicos configurados con LNC=2.
• Las dependencias de software incluyen PyTorch 2.14.0, Transformers 5.16.0 y torchtitan-neuron 0.2.2+neuron.1.
• La carga de checkpoints de Hugging Face verifica la corrección comprobando que la pérdida del paso 1 esté cerca de 6.54 para el modelo Qwen3-4B-Instruct-2507.
• Los kernels NKI personalizados deben cumplir con restricciones específicas de longitud de secuencia, como ser múltiplos de 128 * LNC.
• El paralelismo de tensores fragmenta los pesos de las capas entre rangos (ranks), mientras que el paralelismo de datos fragmenta los parámetros y los estados del optimizador.
• El registro de anulación previene conflictos deteniendo la ejecución si dos anulaciones reclaman el mismo nodo y advirtiendo si ningún nodo coincide.

Por qué es importante

Para ingenieros de software y profesionales de ML, esta guía reduce la barrera de entrada para entrenar modelos grandes en hardware no-NVIDIA. Al mantener la compatibilidad con interfaces estándar de PyTorch y ecosistemas de Hugging Face, disminuye la necesidad de refactorizar extensivamente el código. Los equipos pueden aprovechar repositorios de modelos y formatos de checkpoint existentes, minimizando la sobrecarga operativa asociada con la adopción de nuevas arquitecturas de aceleradores.

La capacidad de inyectar kernels NKI personalizados es particularmente significativa para la optimización del rendimiento. A medida que los modelos crecen en complejidad, las implementaciones genéricas pueden no utilizar plenamente las capacidades del hardware. Permitir a los desarrolladores reemplazar operaciones específicas, como mecanismos de atención o embeddings rotacionales, con kernels ajustados manualmente habilita un control más fino sobre la eficiencia computacional. Esta flexibilidad es crucial para lograr un throughput competitivo y rentabilidad en trabajos de entrenamiento a gran escala.

Además, los pasos detallados de verificación brindan confianza en la fiabilidad del proceso de entrenamiento. Confirmar que los checkpoints exportados se carguen correctamente asegura que los modelos entrenados en Trainium puedan desplegarse sin problemas en entornos de producción que dependen de stacks de inferencia estándar. Esta validación de extremo a extremo respalda prácticas robustas de MLOps y facilita la colaboración entre equipos que utilizan diferentes backends de hardware.

Qué puedes hacer

• Configurar una instancia trn2.48xlarge con la AWS Deep Learning AMI para Neuron o un contenedor Neuron PyTorch.
• Instalar las versiones especificadas de torchtitan-neuron y dependencias utilizando el archivo requirements.txt proporcionado.
• Descargar los artefactos del modelo Qwen3 desde Hugging Face Hub a un directorio local para el entrenamiento.
• Configurar el script de entrenamiento para usar fragmentación de paralelismo de datos en todos los 64 rangos estableciendo data_parallel_shard_degree = -1.
• Verificar las salidas de kernels personalizados contra implementaciones base utilizando entradas float32 para asegurar la precisión numérica.
• Monitorear el uso de memoria de alto ancho de banda por rango para ajustar adecuadamente los tamaños de lote y longitudes de secuencia.