Dos bugs ocultos a simple vista: Una historia de detectives de depuración en vLLM — Asaf Gardin y Yuval Belfer
AAI Engineer
Computing/SoftwareInternet Technology
Transcript
00:00:00Así que pones tu modelo en alguna parte, pones tu agente, cruzas los dedos,
00:00:24esperas que algo falle, todo se ve bien, todo está en orden.
00:00:29Y ves esto.
00:00:33Y ese es el problema, ¿verdad?, con este tipo de errores.
00:00:37No hay fallas, no hay advertencias, ningún error, y la confianza es alta.
00:00:43Eso no es un problema de calidad.
00:00:46Porque, claro, es algo donde realmente no sabes el qué ni el porqué.
00:00:53Bienvenidos a esta charla. Mi nombre es Yuval. Él es Asaf.
00:00:58Juntos, los llevaremos a través de cómo terminamos solucionando este tipo de errores.
00:01:04Un poco sobre nosotros.
00:01:06Trabajamos en AI21, que es un laboratorio de investigación en IA.
00:01:10Empezamos como una empresa de modelos fundacionales, más conocidos por Jamba, que es una arquitectura híbrida entre transformers y Mamba, un estado SSM.
00:01:22Y mientras hacíamos eso, mientras entrenábamos esos modelos, mientras los llevábamos a producción, teníamos usuarios y una carga de trabajo, nos topamos con varios errores interesantes.
00:01:36Y creo que estos son los errores más difíciles de manejar.
00:01:41Porque no es un problema de calidad.
00:01:43No es algo en lo que puedas poner a tu equipo de investigación a optimizar, resolver o hacer que el modelo sea mejor en algo.
00:01:53Así que esto es un problema de ingeniería. Es un caso donde hay alta confianza, pero el resultado es malo.
00:02:02Profundicemos en el primer caso, lo que llamamos la solicitud impostora, para poner en contexto de qué estamos hablando.
00:02:12Hablamos de cómo nosotros, durante el entrenamiento de nuestro modelo Jamba, específicamente hicimos GRPO, un tipo de entrenamiento por RL.
00:02:22Y de nuevo, este es un modelo híbrido.
00:02:24Capas de Mamba y atención.
00:02:27Y para asegurarnos de que estamos alineados, ¿cómo es el flujo de una solicitud?
00:02:31Veamos la vida de una solicitud.
00:02:34Empezamos con el prompt, la tokenización, y luego en la pasada hacia adelante hacemos tanto el prefill como la decodificación.
00:02:44Tras eso, terminamos la pasada hacia adelante y la detokenización para volver al texto.
00:02:51Y lo grave de esto es que falla a muchos niveles, pero principalmente en estos tres.
00:02:59Es lo que llamamos el galimatías uno en mil.
00:03:01No es algo que ocurra en las primeras 500 o 900 solicitudes, sino en la número mil.
00:03:09Lo cual es lo suficientemente raro como para duplicarlo con facilidad, pero demasiado común para lanzarlo.
00:03:18Perdón.
00:03:19Además, solo ocurría en VLLM, no en ningún otro marco de inferencia.
00:03:26Y es de aparición tardía.
00:03:28No es algo que pase si solo tienes unas pocas solicitudes.
00:03:32Necesitas cierto nivel de carga.
00:03:34Así que es raro.
00:03:35Es de aparición tardía.
00:03:36Y es muy específico del motor.
00:03:38Es una tarea muy, muy difícil, y tuvimos que traer a uno de nuestros mejores detectives para encargarse.
00:03:43Así que le cederé la palabra a Asaf para que explique cómo.
00:03:47De acuerdo.
00:03:48Hola a todos.
00:03:48Gracias, Yuval.
00:03:49Gracias, Yuval.
00:03:50Bueno, empezaremos intentando reproducir algo que era muy difícil de reproducir.
00:03:55Básicamente, VLLM tiene muchos parámetros, muchos flags en la CLI y muchas opciones que se pueden ajustar y modificar.
00:04:04Y una de las cosas que nos ayudó a entender cómo reproducirlo, porque al principio al intentar replicarlo solo enviando prompts aquí y allá en algunos lotes, no lográbamos que el modelo devolviera galimatías.
00:04:17El modelo respondía perfectamente.
00:04:19Así que lo que hicimos fue intentar provocarlo en muy, muy poco tiempo.
00:04:25Para tener un bucle de retroalimentación rápido al depurarlo.
00:04:29Entonces, lo que hicimos fue tomar uno de los parámetros más predeterminados y comunes de VLLM, que es la utilización de memoria de la GPU, el cual te permite elegir cuánta memoria de GPU deseas asignar para los pesos, activaciones, caché KV, etcétera.
00:04:48Y la redujimos del 90% al 20%.
00:04:51Y una vez que hicimos eso y empezamos a ejecutar muchas solicitudes simultáneamente, de repente la solicitud número, digamos, 854 devolvió galimatías.
00:05:05Y al hacerlo, muestreamos todos los lotes con temperatura cero para poder obtener de forma determinista y constante la misma solicitud respondiendo con galimatías.
00:05:17Así que, como dijo Yuval, solo pasaba en VLLM.
00:05:20Y usamos, para lograrlo y entender de dónde venía realmente el problema, Hugging Face Transformers como referencia, ya que Transformers tenía una implementación muy estándar y básica de nuestros núcleos de Mamba, a diferencia de VLLM, cuyos núcleos y motor han sufrido muchos cambios para admitir funciones avanzadas.
00:05:46Así que usamos Transformers como línea base para entender si había o no un problema con nuestra inferencia o con el modelo.
00:05:55Lo que hicimos fue tomar VLLM, enviar todos nuestros prompts a través de él y generar una respuesta, todas las respuestas.
00:06:02Ahora tenemos en nuestras manos la respuesta junto con las probabilidades logarítmicas, ya que en VLLM puedes extraerlas e inspeccionarlas.
00:06:11Luego tomamos la secuencia completa, el prompt y la generación, y se la pasamos a la pasada hacia adelante de Hugging Face.
00:06:20Pero solo ejecutamos el prefill, obtuvimos los logits de la respuesta del prefill, los pasamos por Softmax y así pudimos comparar la divergencia en las distribuciones de nuestros tokens.
00:06:38Aquí tienen un pequeño pseudocódigo de cómo se veía eso.
00:06:41Pueden ver que tomamos el prompt, lo ejecutamos mediante el generate de VLLM, obtenemos la respuesta con las log probs, lo pasamos a la pasada hacia adelante de Hugging Face y solo ejecutamos el prefill.
00:06:55Creamos una función llamada compute log probs, que ejecuta Softmax, calcula la diferencia y permite ver la divergencia entre las log probs de cada token.
00:07:09Bien, ahora que tenemos las herramientas para entender de dónde podría venir el problema, empezamos a buscar diferentes sospechosos en el motor de VLLM.
00:07:19Lo primero que revisamos fue el kernel de prefill CUDA de Mamba.
00:07:24Lo examinamos, inspeccionamos todas las operaciones matemáticas y revisamos los tensores de entrada y salida antes y después de llamar al prefill.
00:07:36Todo se veía bien.
00:07:38Lo segundo que hicimos fue ejecutar la herramienta compute sanitizer de NVIDIA para ver si teníamos desbordamientos de memoria u otros errores de memoria.
00:07:48Parecía estar bien.
00:07:52Luego intentamos aislar entre los kernels de decodificación y los de prefill.
00:07:56Vimos que los de prefill funcionaban correctamente, así que intentamos no llamar a los de decodificación, ya que en Mamba es posible hacerlo.
00:08:03Lo que hicimos fue redirigir todas nuestras llamadas y cálculos a través del kernel de prefill, y ahí lo tienen.
00:08:12El galimatías de repente desapareció.
00:08:14Así que dijimos: de acuerdo, tienen que ser los kernels de decodificación.
00:08:18Pero ya saben cómo es esto en el software.
00:08:20Te emociona demasiado pronto y luego descubres que no era eso.
00:08:24Así que empezamos a jugar con el motor de VLLM, tuvimos que levantar el capó y ver qué podíamos hacer para comprenderlo mejor y ensuciarnos las manos.
00:08:38Porque VLLM no nos daba más herramientas para depurar a fondo nuestro kernel y nuestra pasada hacia adelante.
00:08:45Una vez que un tensor o solicitud llega a la pasada hacia adelante, antes de entrar a los kernels de prefill y decodificación, no tienen una identidad clara.
00:08:55No se puede saber realmente qué prompt se está procesando.
00:08:59Todo son solo tensores, números y matrices.
00:09:02Lo que hicimos fue añadir el ID de solicitud a una clase llamada forward context y propagarlo hasta la pasada de Mamba, justo antes de llamar a los kernels de prefill y decodificación.
00:09:16Y allí pudimos poner una simple condición if con el ID de la solicitud que daba galimatías y colocar un punto de interrupción.
00:09:26De ese modo pudimos deducir e inspeccionar todos los metadatos asociados.
00:09:33Y en cuanto hicimos eso, vimos que por primera vez al pasar por la pasada hacia adelante, la solicitud estaba realizando decodificación antes que el prefill.
00:09:43El planificador decidió que esta solicitud debía hacer decodificación antes del prefill.
00:09:50Y como dijo Yuval antes, en el ciclo de vida de un prompt, este debe pasar primero por el prefill y luego por la decodificación.
00:10:00Lo que ocurría era que al ejecutar una solicitud en Mamba con decodificación primero, después de que ya se habían calculado muchas otras, el estado ya estaba sobreutilizado.
00:10:15Y estábamos usando los datos y cálculos de solicitudes obsoletas que habían venido antes.
00:10:21Así que terminábamos ejecutando la decodificación sobre solicitudes anteriores.
00:10:25Y eso generaba galimatías para nosotros.
00:10:28Así que los kernels no estaban haciendo algo incorrecto.
00:10:30Se les llamaba en el momento equivocado para las solicitudes equivocadas.
00:10:34¿Y por qué importaba solo para Mamba?
00:10:37La razón era que en attention, escribes los tokens KV antes de leerlos.
00:10:47Así que aunque tengas datos obsoletos, se sobrescriben.
00:10:50Pero en Mamba, como mencioné, al pasar primero por los kernels de decodificación, lees el estado y luego calculas sobre él.
00:10:58Por lo tanto, terminabas usando datos obsoletos al realizar la decodificación.
00:11:05Y la solución fue relativamente sencilla.
00:11:07Solo debíamos asegurarnos de que cuando el planificador clasifica una solicitud por primera vez, si ve una cuyos tokens nunca se han calculado y son cero,
00:11:26los marque como prefill para que al llegar a la pasada hacia adelante se usen para prefill y no para decodificación ni de forma fragmentada.
00:11:38Pueden ver que se integró al código poco después.
00:11:41Y eso nos lleva —y entonces pensamos que todo estaba resuelto, ¿verdad?
00:11:45Creíamos que todo estaba solucionado y listo.
00:11:47No más problemas.
00:11:48Pero eso fue casi así, porque poco después nos condujo al caso número dos,
00:11:54el cual sacó a la luz otro inconveniente que enfrentamos en nuestro RL y en la inferencia.
00:12:01Así que ejecutamos RL, y mientras realizábamos nuestros entrenamientos, nuestros postentrenamientos, y analizábamos nuestras evaluaciones y todos nuestros benchmarks,
00:12:10creíamos que teníamos algunos picos de logprop entre el rollout y el paso de FSDP.
00:12:15Y eso era antes de cualquier actualización de pesos.
00:12:18Así que mismos pesos, mismas entradas, y los dos logprops deberían ser idénticos.
00:12:23Pues no lo eran.
00:12:24Vimos que constantemente cada 12 pasos había un pico de logprop, y eso era un poco extraño.
00:12:33Ahora bien, ¿qué harían ustedes, verdad?
00:12:36¿Qué es lo primero que hay que hacer aquí?
00:12:39Así que queríamos encontrar alguna palanca que cambiara cómo fallan las cosas y no solo cuánto fallan.
00:12:44Queremos ver cuánto... queremos ajustar algunos parámetros que no solo nos indiquen, oye, este error... este error es muy, muy grave... este error ocurre tantas veces, o... y así sucesivamente.
00:12:57Y por eso queríamos ajustar algunos parámetros que nos indicaran que, una vez ajustados, comprendemos cómo se conectan con cualquier elemento del motor de VLLM.
00:13:08Y así podremos ir específicamente a depurar esa parte en concreto.
00:13:14Ese es un meme divertido que todos querían incluir.
00:13:19Así que lo que hicimos fue decidir aumentar la cantidad de rollouts por prompt.
00:13:26Dado que en nuestro motor de RL predeterminado tenemos ocho rollouts por prompt, y vimos que ocurría de forma determinista cada 12 pasos, decidimos: de acuerdo, intentemos subirlo un poco y aumentar la cantidad de rollouts por prompt.
00:13:40Así que empezamos a duplicarlo de 8 a 16, a 64, 32 y 128.
00:13:45Y aquí pueden ver que es casi... casi... hay un patrón aquí.
00:13:51Cuanto más lo aumentábamos, más rápido ocurría, porque lo que queríamos lograr aquí era intentar reproducir el problema lo más rápido posible para tener un bucle de depuración y de retroalimentación más rápido.
00:14:05Así que cuando lo ejecutamos con 128 rollouts por prompt, ocurrió inmediatamente en el primer paso, y no tuvimos que esperar al paso 12, 24 y así sucesivamente.
00:14:12Ahora bien, podrían pensar, de acuerdo, así que antes jugaron con la utilización de memoria de la GPU.
00:14:18La ajustaron, la redujeron.
00:14:20Parece que, ya saben, al someterla a presión, realmente saca a la luz los problemas.
00:14:24Así que nosotros también pensamos eso.
00:14:26Y cuando redujimos la memoria de la GPU de 0.9 a 0.2, en realidad el problema desapareció.
00:14:35Así que tiramos de la palanca equivocada.
00:14:38Y la razón es porque notamos que los núcleos de Mamba utilizaban un puntero de índice entero sin signo de 32 bits.
00:14:48Así que una vez que el desplazamiento superaba unos, ya saben, 4 mil millones de números, volvía a empezar desde cero en lugar de arrojar un error.
00:14:55Por lo tanto, cuando redujimos la memoria de la GPU, VLLM asignó un búfer de estado pequeño, y el índice de caché nunca llegó a ser lo suficientemente grande como para alcanzar ese límite.
00:15:04Así que simplemente no estábamos avanzando lo suficiente como para que el búfer provocara un desbordamiento.
00:15:09De nuevo, la solución fue bastante simple.
00:15:13Todo lo que teníamos que hacer era cambiar una sola palabra, una variable de tipo de dato de UINT32 a size_t, lo que básicamente significa que para la mayoría de las arquitecturas modernas de hardware,
00:15:27size_t pasaría a significar que ahora se cambia a un entero sin signo de 64 bits.
00:15:32Y ese es un número muy grande.
00:15:34Nunca alcanzamos ese número y ese desbordamiento ya no volvió a ocurrir.
00:15:38Así que lo que podemos ver aquí es que tuvimos dos escenas y un solo culpable.
00:15:45Ambos, digamos, tenían síntomas similares.
00:15:48Ambos presentaban jerga silenciosa y picos silenciosos de log prob, que también a veces generaban jerga.
00:15:55Ambos estaban relacionados con la caché de estado de Mamba.
00:15:57Ambos salieron a la luz debido a la presión de la memoria, ya fuera para bien o para mal.
00:16:03Y ambos fueron encontrados mediante el análisis forense de log prob.
00:16:09Los sistemas de inferencia con estado no fallan ruidosamente.
00:16:12Te mienten con total seguridad.
00:16:13O sea, obviamente, a veces hay cuelgues.
00:16:15Obtienes errores de fuera de límites.
00:16:16Obtienes otras excepciones, ya sabes, y todo eso.
00:16:19Pero a veces hay errores que no salen a la superficie y no obtienes un registro de seguimiento.
00:16:24No obtienes nada.
00:16:25Tienes que investigar y comprender por qué ocurren las cosas.
00:16:29Así que, si hay alguna conclusión que sacar de esta presentación, es crear el script de comparación de log probs.
00:16:36Si necesitas comparar tu calidad, debes compararla para entender si modularizaste los problemas o no.
00:16:41Un script de comparación de log probs con una línea base de algún otro framework de inferencia que tengas o hayas construido siempre es genial.
00:16:48Reproducir bajo presión y memoria limitada, subir la escala al máximo, jugar con otros parámetros que te ofrece el framework de inferencia e intentar comprender realmente de dónde viene el problema.
00:16:59Busca qué es lo que altera la forma del fallo, el tiempo, el espacio y la ubicación.
00:17:05Y cuando las cosas no tengan identidad, propaga la identidad a través de ellas.
00:17:10Y lo que también quiero que se lleven de esto, y no teman incluso, ya saben, en sistemas complejos como VLLM o cualquier otro framework complejo, no teman meterse en el código, mancharse las manos.
00:17:23A veces, ya saben, los modelos de lenguaje, los LLMs, tal vez te digan cómo funcionan las cosas, pero sin verlo con tus propios ojos, sin mancharse las manos, no obtendrás una comprensión completa de lo que está pasando.
00:17:34Gracias, pueden agregarnos en LinkedIn, escaneen el código QR para leer el blog real que publicamos con este descubrimiento.
00:17:43Sí, eso es todo.
00:18:04Nos vemos la próxima vez.
Community Posts
No posts yet. Be the first to write about this video!
Write about this video