Hoy, Amazon anunció que SGLang-Neuron está disponible públicamente y se puede acceder a él a través de un contenedor de aprendizaje profundo de AWS, lanzado por RadixArk en colaboración con Annapurna Labs. Con SGLang-Neuron, puede servir modelos y generar implementaciones de aprendizaje por refuerzo (RL) con SGLang en AWS Trainium utilizando la misma interfaz SGLang que utiliza hoy. Trainium es una cartera de chips de IA de AWS que ofrece una sólida economía para el entrenamiento y la inferencia de IA a escala. Ahora, ejecutar SGLang en Trainium sigue el flujo de trabajo de SGLang que ya conoce: inicia modelos a través de la interfaz del servidor de SGLang e invoca el punto final utilizando las API de la misma generación. Debajo de esa interfaz, SGLang-Neuron se ejecuta en PyTorch nativo. Con su diseño de complemento fuera del árbol, la compatibilidad con Trainium y el núcleo SGLang se desarrollan de forma independiente y se iteran más rápido.
Cómo se ejecuta SGLang en Trainium
SGLang-Neuron es un complemento fuera del árbol que conecta SGLang con AWS Trainium. La instalación del complemento registra la plataforma neuronal a través de los puntos de entrada del complemento de SGLang, y la configuración SGLANG_PLATFORM=neuron la selecciona cuando se inicia el servidor. SGLang core continúa proporcionando el servidor API, el ciclo de vida de las solicitudes, RadixAttention y la contabilidad de caché de prefijos. El complemento implementa los modelos específicos de Trainium, la atención, la combinación de expertos (MoE) y el corredor de gráficos.
Figura 1: SGLang-Neuron en AWS Trainium
TorchNeuron registra Trainium como un dispositivo PyTorch nativo. SGLang-Neuron utiliza operaciones estándar de PyTorch para la ruta de ejecución del modelo e incorpora núcleos escritos con Neuron Kernel Interface (NKI) para mejorar el rendimiento. Compila estas operaciones y núcleos en programas ejecutables utilizando torch.compile, que el corredor de gráficos prepara y ejecuta durante el servicio.
SGLang-Neuron también admite la desagregación de PD. En esta configuración, un servidor de precarga calcula el caché KV del mensaje y lo transfiere a un servidor de decodificación usando NIXL sobre Elastic Fabric Adapter (EFA). Luego, el servidor de decodificación genera los tokens restantes. Cada servidor todavía utiliza el programador de SGLang y el ciclo de vida de las solicitudes, mientras que el complemento administra la ejecución de Trainium y la ruta de transferencia de caché KV.
Sirviendo modelos en AWS Trainium usando SGLang-Neuron
Para seguir este tutorial, utilice una instancia de Trainium3 Amazon Elastic Compute Cloud (Amazon EC2) con un contenedor pytorch-inference-sglang-neuronx. Aquí están los pasos:
Inicie el contenedor.
Inicie el servidor del modelo SGLang dentro del contenedor.
Envíe una solicitud al servidor modelo.
Para obtener información adicional, consulte la guía de introducción de SGLang-Neuron para obtener instrucciones de configuración y explore la documentación del repositorio para conocer los modelos compatibles, las funciones y las configuraciones validadas.
Conclusión
SGLang-Neuron trae el flujo de trabajo de servicio familiar de SGLang a AWS Trainium a través de PyTorch nativo. Puede comenzar iniciando el contenedor SGLang-Neuron y siguiendo el inicio rápido para ofrecer un modelo compatible. Pruébelo con sus cargas de trabajo, contribuya a través del repositorio de proyectos y comparta comentarios para ayudar a dar forma a futuros modelos y compatibilidad con funciones.
Sobre los autores
Alex Nails
Yihao Wang
Lingyan Hao
Selvam Srecharan
Zeyu Wang
Vinesh Ravuri
Yuchen Jiang
Panigrahi rosa