El crecimiento exponencial de los modelos de Machine Learning, especialmente los Large Language Models (LLMs), ha exacerbado la brecha entre la capacidad de memoria necesaria para alojar los pesos del modelo y el costo/disponibilidad de High Bandwidth Memory (HBM) o DRAM. Este problema fundamental de la computación, que se manifiesta como una restricción de capacidad de memoria en el dispositivo, limita el tamaño de los modelos que pueden ejecutarse eficientemente en un solo acelerador y fuerza estrategias de sharding complejas que introducen latencia por comunicación entre dispositivos.
HBF emerge como una propuesta para mitigar esta limitación. Al integrar memoria flash de alta densidad en el mismo paquete que el chip de cómputo, HBF busca ofrecer una capacidad significativamente mayor que HBM a un costo por bit inferior, manteniendo al mismo tiempo un ancho de banda superior al de los SSDs externos. Sin embargo, esta solución no es una extensión transparente de la memoria, sino que introduce un nuevo nivel en la jerarquía de memoria con características de acceso más cercanas a las de un dispositivo de almacenamiento masivo, requiriendo adaptaciones profundas en el software para su aprovechamiento.
La relevancia de HBF radica en su potencial para desbloquear la ejecución de modelos más grandes en configuraciones de hardware existentes, o para reducir la necesidad de sharding agresivo, lo que a su vez podría simplificar la arquitectura de sistemas distribuidos para ML. Su éxito dependerá de la capacidad de los frameworks de ML para abstraer la complejidad de su interfaz de acceso y de si los beneficios de capacidad superan los desafíos de rendimiento inherentes a su menor ancho de banda y latencia de acceso comparada con HBM.
Arquitectura del Sistema
HBF se concibe como un conjunto de cubos de memoria flash apilados en el mismo paquete que el chip de cómputo, adyacente a HBM. A nivel de hardware, se diferencia de HBM en que no es una extensión de la memoria principal con acceso de byte, sino un dispositivo de almacenamiento de bloques. La interacción con HBF se realiza mediante operaciones de DMA (Direct Memory Access), donde el software es responsable de mover bloques de datos entre HBF y DRAM/HBM. Esto implica que el software debe gestionar explícitamente la localidad de los datos y el ciclo de vida de los bloques, similar a cómo un controlador de SSD gestiona la memoria flash.
Las operaciones de acceso a HBF deben ser en bloques grandes y alineados, no a nivel de byte. Esto significa que una modificación de un solo byte podría requerir la lectura de un bloque completo (ej. 64 KB) en DRAM, la modificación del byte en DRAM, y la escritura del bloque completo de vuelta a HBF. Esta característica es fundamentalmente diferente del acceso a memoria volátil y requiere que el software implemente lógicas de caching y buffering para optimizar el rendimiento y la vida útil de la flash (wear leveling). La gestión de la retención de datos y el wear leveling, tradicionalmente manejados por el firmware del controlador de SSD, recae parcialmente en el software de la aplicación o en un runtime especializado.
Para cargas de trabajo de ML, HBF se propone para almacenar componentes de modelos que no requieren acceso de baja latencia o alta frecuencia, pero sí una gran capacidad. Ejemplos incluyen expertos de Mixture-of-Experts (MoE) que se cargan bajo demanda, o partes "frías" de la caché KV en arquitecturas de atención dispersa. En estos escenarios, el software orquesta el DMA de los bloques de datos relevantes desde HBF a HBM/DRAM justo antes de que sean necesarios para el cómputo, y potencialmente los descarga de nuevo a HBF cuando ya no están "calientes". Esto reduce la presión sobre la capacidad de HBM/DRAM y la comunicación entre dispositivos, ya que más pesos pueden ser replicados localmente en HBF.
Flujo de Carga de Expertos MoE con HBF
- 1 Modelo ML Identifica experto MoE necesario para el cómputo actual.
- 2 Runtime ML Verifica si el experto está en HBM/DRAM. Si no, inicia carga.
- 3 Software HBF Calcula dirección de bloque en HBF para el experto.
- 4 DMA Engine Transfiere bloque de datos del experto desde HBF a HBM/DRAM.
- 5 HBM/DRAM Experto disponible para el cómputo en memoria de alta velocidad.
- 6 Modelo ML Ejecuta cómputo con el experto cargado.
| Capa | Tecnología | Justificación |
|---|---|---|
| storage | High Bandwidth Flash (HBF) | Proporcionar almacenamiento de alta capacidad y ancho de banda moderado, integrado en el paquete del procesador, para componentes de modelos ML que no caben en HBM/DRAM. vs HBM (High Bandwidth Memory), DRAM (Dynamic Random-Access Memory), SSD (Solid State Drive) externo |
| compute | GPU/CPU | Ejecutar las cargas de trabajo de Machine Learning, utilizando HBF como una extensión de la jerarquía de memoria para acceder a modelos más grandes. |
| orchestration | Custom ML Runtime | Gestionar el movimiento de datos entre HBF y HBM/DRAM, implementando lógicas de caching, prefetching y gestión de bloques para optimizar el uso de HBF. vs Modificaciones del kernel del SO, Frameworks de ML existentes (ej. vLLM sin modificaciones) |
Trade-offs
Ganancias
- ▲▲ Capacidad de memoria en el paquete
- ▲ Costo por capacidad
- ▲ Reducción de comunicación cross-device
Costes
- ▲ Ancho de banda por unidad de costo (vs HBM)
- ▲ Latencia de acceso (vs HBM/DRAM)
- ▲▲ Complejidad del software
- ▲ Granularidad de acceso
Fundamentos Teóricos
El concepto de HBF se alinea con la investigación en jerarquías de memoria no volátil y la gestión de memoria heterogénea, un campo explorado extensamente desde los primeros días de la computación. La idea de tratar la memoria flash como un nivel en la jerarquía de memoria, en lugar de solo como almacenamiento persistente, ha sido objeto de numerosos estudios. Trabajos como los de Intel Optane (3D XPoint) intentaron cerrar la brecha entre DRAM y NAND flash, ofreciendo persistencia y acceso a nivel de byte, aunque con latencias superiores a DRAM.
La necesidad de gestionar accesos en bloques grandes y alineados en HBF, y la delegación de funciones de controlador de almacenamiento al software, resuena con los principios de sistemas de archivos y bases de datos que operan directamente sobre dispositivos de bloques. Conceptos como el Write-Ahead Log (WAL) y las estructuras de datos basadas en LSM-tree (Log-Structured Merge-tree), populares en bases de datos NoSQL, fueron diseñados para optimizar escrituras secuenciales en medios de almacenamiento con características similares a la flash (borrado de bloques, escrituras fuera de lugar). La gestión de la caché KV en HBF, por ejemplo, podría beneficiarse de principios de locality-aware caching y prefetching, temas clásicos en la optimización de sistemas de memoria y almacenamiento. La gestión de la memoria flash a bajo nivel, incluyendo el wear leveling y la recolección de basura, ha sido un área activa de investigación en sistemas operativos y sistemas de archivos, con papers que datan de las primeras implementaciones de flash en la década de 1990.