La migración del modelo Gemma-4 31B de Google a un AWS Inferentia2 inf2.24xlarge produjo una coincidencia perfecta token por token con la referencia de CPU; sin embargo, cada frase generada era un sinsentido. La brecha entre "coincidir" y "funcionar" sirve ahora como advertencia para cualquiera que intente exprimir LLMs masivos en los chips de inferencia personalizados de Amazon.

Por qué una coincidencia token por token no es suficiente

El desarrollador comparó cada token de salida del dispositivo Inferentia con el token producido por una ejecución del modelo en CPU. Los flujos eran idénticos, por lo que el hardware parecía haber reproducido exactamente la implementación de referencia. En realidad, ambos flujos introdujeron un prompt mal formado en un modelo al que se le había quitado su plantilla de chat (chat template) y se le habían suministrado marcadores de turno incorrectos. La falta de la plantilla envió al modelo a un bucle infinito, generando disparates. El hardware cumplió su función: reprodujo un error que ya existía en el código de referencia.

La lección es sencilla: SEQ_MATCH (igualdad secuencial de tokens) no es sinónimo de corrección. Si la implementación de referencia está rota, una réplica de hardware fiel heredará el mismo fallo. La validación debe ir más allá de la paridad a nivel de token; requiere comprobaciones funcionales de extremo a extremo con entradas correctamente formateadas.

Buffers que se hacen pasar por parámetros

Durante la fase de carga, el cargador del modelo omitió un componente llamado layer_scalar. El código registró este objeto como un buffer en lugar de un parámetro en la definición del modelo de PyTorch. Los buffers son tensores estáticos que el entrenamiento no actualiza, y muchos cargadores los ignoran al convertir a formatos compatibles con Neuron. Al omitirlo, los factores de escala de varias capas se quedaron en sus valores por defecto, distorsionando los cálculos en toda la red. No se produjo ningún error; el modelo se compiló y el pipeline de inferencia se ejecutó, pero los resultados numéricos eran incorrectos.

Para cualquiera que esté migrando modelos grandes a Inferentia, audite cada tensor que no sea un parámetro. Incluso si un tensor no está destinado a ser aprendido, puede seguir siendo esencial para el cálculo correcto del paso hacia adelante (forward-pass). Verificar manualmente la inclusión de los buffers puede prevenir errores de escala silenciosos que, de otro modo, serían difíciles de diagnosticar.

La volatilidad de las instancias spot y la compilación de 39 minutos

Ejecutar un modelo de 31 mil millones de parámetros en una instancia spot parece económico, pero el ahorro conlleva eventos de recuperación impredecibles. El tiempo de compilación del desarrollador —unos 39 minutos para traducir el modelo a código compatible con Neuron— se esfumó cuando AWS reclamó la instancia. Para sobrevivir a las interrupciones, construyeron una red de seguridad de tres pilares:

  • ModelBuilder mantuvo el uso de memoria dentro del límite de 384 GB del host, evitando bloqueos que obligarían a un reinicio.
  • El espejo inmediato en S3 tanto de los archivos de pesos originales como de los "neffs" compilados (archivos ejecutables de Neuron) permitió que una nueva instancia retomara el trabajo exactamente donde la anterior lo había dejado.
  • Un poller multi-región escaneó las regiones de AWS en busca de capacidad spot disponible y lanzó una nueva instancia tan pronto como apareció una.

Estos pasos convirtieron una compilación frágil de punto único en un pipeline resiliente que sobrevive a la volatilidad de los mercados spot.

Trampas del sharding con diseños de atención mixtos

Gemma-4 31B utiliza dos configuraciones de atención. Algunas capas emplean cuatro cabezales de clave-valor (KV), otras un número diferente. Dividir el modelo uniformemente en ocho rangos paralelos falla cuando el número de cabezales KV de una capa no es divisible de forma exacta. Intentar realizar el sharding de una capa de 4 cabezales en ocho rangos obligaría a cada rango a manejar medio cabezal, una imposibilidad matemática que provoca desajustes de forma (shape mismatches) y errores en tiempo de ejecución.

La solución fue replicar las capas con sharding global (aquellas con recuentos de cabezales compatibles) en todos los rangos y solo realizar el sharding de las capas "deslizantes" cuyos recuentos de cabezales permitieran una división equitativa. Esta estrategia híbrida mantuvo la eficiencia del paralelismo de tensores (tensor-parallel) al tiempo que evitaba la división ilegal de los cabezales KV, eliminando los errores de paralelización de tensores que afectaron a los intentos anteriores.

Conclusión

Migrar un LLM gigante a Inferentia es más que un ejercicio de "compilar y ejecutar". Exige pruebas funcionales rigurosas más allá de la igualdad de tokens, una verificación meticulosa de que cada tensor —ya sea parámetro o buffer— se maneje correctamente, y una estrategia de despliegue que anticipe la recuperación de instancias spot. Por último, el sharding debe respetar la geometría de atención interna del modelo; de lo contrario, el paralelismo que promete velocidad se convierte en una fuente de fallos silenciosos.