- Se optimizó GPT-OSS-120B, el LLM de código abierto de OpenAI, para lograr un rendimiento de más de 500 tokens por segundo en entornos de GPU de NVIDIA.
- Se realizaron pruebas paralelas de diversos frameworks de inferencia como TensorRT-LLM, vLLM y SGLang, con soporte para las arquitecturas Hopper y Blackwell.
- Se corrigieron bugs de compatibilidad e integró un formato de respuesta nuevo como Harmony, además de aplicar optimizaciones como enrutamiento con conocimiento del caché KV y decodificación especulativa basada en Eagle.
- Al comparar Tensor Parallelism y Expert Parallelism, se eligió Tensor Parallelism por su menor latencia, y en Blackwell se utilizó el backend MoE de TensorRT-LLM.
- Se planea seguir optimizando el rendimiento, incluyendo la decodificación especulativa (Speculative) con un modelo de “draft” más pequeño.
Visión general
- Cuando OpenAI lanzó GPT-OSS-120B, su último modelo de lenguaje de gran escala de código abierto, Baseten se propuso alcanzar el máximo rendimiento.
- Baseten es el socio de lanzamiento oficial de OpenAI.
- Con datos públicos de usuarios reales de OpenRouter, validaron un rendimiento líder en entornos con GPU de NVIDIA frente a otros proveedores.
- Gracias al Flexible Inference Stack y la experiencia del equipo de ingeniería de modelos, aplicaron parches de optimización de forma rápida, casi por horas.
- En las pocas horas de redactar el blog, lograron un aumento adicional de más de 100 tokens por segundo y mantuvieron un 100% de tiempo de actividad.
Esfuerzos de optimización de rendimiento
- Se realizaron pruebas y benchmarking con distintos frameworks de inferencia como TensorRT-LLM, vLLM y SGLang.
- En paralelo, se aseguró compatibilidad con las arquitecturas de GPU Hopper y Blackwell.
- Se integró con componentes clave como el Flexible Inference Stack de Baseten y NVIDIA Dynamo.
- Se aplicaron técnicas de optimización de rendimiento ampliamente validadas, como KV cache-aware routing y Speculative decoding (basado en Eagle).
A continuación, se detallan los pasos clave para lograr simultáneamente un rendimiento SOTA y soporte de ventana de contexto completa.
Paso 1: Ejecución inicial de inferencia
- Cualquiera sea el enfoque, el punto de partida fue ejecutar la inferencia inicial (baseline inference) lo más rápido posible.
- Aprovechando la potencia de la GPU, varios ingenieros corrieron en paralelo experimentos con vLLM, SGLang y TensorRT-LLM.
- Se logró poner en marcha rápidamente el TensorRT-LLM, que mostró el mejor rendimiento.
- Se garantizó soporte de TensorRT-LLM tanto en Hopper (la de mayor cantidad de GPUs H100) como en Blackwell (donde B200 ofrece mejor velocidad).
- Gracias a la flexibilidad de Baseten Inference Runtime, resultó sencillo adaptar nuevos modelos de arquitectura y cambiar herramientas dentro del stack rápidamente.
Paso 2: Corrección de bugs de compatibilidad
- La aparición de nuevas arquitecturas de modelos suele venir acompañada de fallos frecuentes durante la integración de frameworks.
- GPT-OSS agregó nuevas tecnologías, como el formato de respuesta Harmony, lo que generó bugs al integrar con frameworks existentes.
- Para mantener velocidad y precisión a la vez, realizaron ajustes y pruebas repetidas, y contribuyeron al open source las correcciones efectivas.
- La colaboración de la comunidad open source global está haciendo que los caminos de optimización y las correcciones de bugs avancen con rapidez.
Paso 3: Optimización de la configuración del modelo
- Aunque OpenAI indica que GPT-OSS-120B funciona en una sola H100, en la práctica se favorece paralelizar entre 4~8 GPUs para mejor rendimiento.
- Tensor Parallelism destaca en latencia, mientras que Expert Parallelism destaca en throughput del sistema.
- Como el objetivo de Baseten era optimizar la latencia, eligieron Tensor Parallelism.
- En Blackwell, se aplicó el TensorRT-LLM MoE Backend, mejorando el rendimiento de los kernels CUDA en comparación con el backend Triton anterior.
- Se publicaron configuraciones optimizadas para entornos Hopper y Blackwell, y en la Model API se adoptó la configuración basada en Blackwell.
Optimización de rendimiento adicional
- Aunque con la primera etapa de optimización ya alcanzaron niveles SOTA en throughput y latencia, aún hay margen para mejorar.
- La actualización principal prevista es la incorporación de Speculative Decoding.
- Este método usa un modelo de “draft” más rápido que genera tokens candidatos, que luego son validados por el modelo principal.
- Baseten recomienda Eagle 3, pero opera de forma dinámica con más de 10 algoritmos dentro del stack de inferencia según el escenario.
- La decodificación especulativa permite inferir varios tokens a la vez para lograr una mejora de velocidad más eficiente
Aún no hay comentarios.