La arquitectura híbrida Falcon-H1 se ha integrado en NVIDIA Megatron Core, un marco fundamental para entrenar modelos de lenguaje a gran escala que ha evolucionado rápidamente. Esta biblioteca de código abierto es conocida por su rendimiento optimizado para GPUs y su paralelismo líder en la industria.
Este artículo ofrece una visión técnica sobre cómo el Technology Innovation Institute (TII), creadores de la familia de modelos Falcon, han contribuido a Megatron Core y los frameworks Megatron Bridge. Se analizará la implementación de la arquitectura híbrida Falcon-H1 y la integración de BitNet en Megatron Core.
Integración de Falcon-H1 en Megatron Bridge
La implementación de la arquitectura híbrida Falcon-H1 resalta los desafíos de coordinar capas heterogéneas de Transformer y Mamba junto con multiplicadores µP no aprendibles. A continuación, se detallan los aspectos de esta integración.
El diseño híbrido se aparta de la disposición secuencial común en otros modelos híbridos recientes. Dentro de cada bloque, el mecanismo de atención y el modelo de estado espacial (SSM) operan en paralelo, concatenando sus salidas antes de proyectarlas en la salida del bloque.
Diseño híbrido paralelo
El núcleo de las contribuciones del TII a Megatron es la arquitectura híbrida paralela Falcon-H1. En lugar de apilar capas distintas, Falcon-H1 adopta un diseño paralelo. Esto permite que los componentes de atención y SSM procesen la entrada simultáneamente dentro de cada bloque de procesamiento.
Las salidas de las ramas de atención y Mamba se concatenan antes de la proyección, fusionando la memoria a largo plazo de los SSM con el modelado de dependencia de largo alcance de la atención. Esta flexibilidad permite explorar distintas configuraciones de capas en el modelo.
Integración de dos repositorios
El soporte de Falcon-H1 abarca dos repositorios con responsabilidades distintas. En Megatron Core, el TII aportó:
- La capa
ParallelHybridLayer, que ejecuta Mamba y atención en paralelo. - La lógica de asignación de capas actualizada que introduce el símbolo
PARALLEL.
Esto también incluye herramientas de conversión de punto de control para cargar y guardar modelos híbridos paralelos.
En Megatron Bridge, se construyó el modelo Falcon-H1 completo sobre estas primitivas, incluyendo componentes adicionales que permiten la conversión de pesos de Hugging Face a Megatron.
