El entrenamiento de modelos de recomendación y ranking presenta un desafío fundamental en la arquitectura de sistemas distribuidos: la comunicación intensiva. A diferencia de los Large Language Models (LLMs) que demandan un throughput masivo de floating-point, los modelos de recomendación están limitados por la necesidad de una comunicación rápida y eficiente entre los aceleradores. Sus tablas de embeddings, que pueden constituir más del 99% de los parámetros del modelo, requieren paralelismo híbrido que genera frecuentes operaciones colectivas como AllReduce, AllToAll y AllGather.
Tradicionalmente, en arquitecturas de GPU, estas operaciones de comunicación compiten por los mismos recursos de cómputo (streaming multiprocessors) que el entrenamiento, llevando a una subutilización del hardware y degradación del rendimiento. MTIA 300 aborda este problema haciendo de la comunicación un ciudadano de primera clase en el diseño del chip, integrando la red y dedicando hardware específico para el manejo de colectivos, lo que permite un desacoplamiento casi perfecto entre cómputo y comunicación.
Arquitectura del Sistema
La arquitectura de MTIA 300 integra dos chiplets de red directamente en el paquete del chip, cada uno con seis NICs RDMA personalizadas de 800 Gbps, proporcionando un ancho de banda I/O total de 1.2 TB/s. Esta integración elimina el cuello de botella host-device-NIC presente en las arquitecturas de GPU tradicionales, donde la CPU intermedia entre el acelerador y la red. Las 12 NICs basadas en Ethernet se utilizan de forma flexible para comunicación scale-up (dentro de un rack, hasta 1 TB/s) y scale-out (entre racks, a 200 GB/s), permitiendo reconfigurar la partición de NICs según las necesidades del modelo.
Además de su matriz de 12x6 elementos de procesamiento (PEs) para cómputo, el chip incluye 16 motores de mensajes (MEs) dedicados que manejan la comunicación de forma independiente. Cada ME contiene un núcleo RISC-V para orquestación, una interfaz NIC para enrutamiento de solicitudes y un bloque de cómputo cercano a la memoria (NMC) que realiza reducciones a 128 bytes/ciclo. Estos NMCs, ubicados junto a la HBM y la caché, ofrecen más de 2.8 TB/s de throughput de reducción, permitiendo la ejecución a velocidad de línea de colectivos AllReduce y ReduceScatter sin involucrar la malla de cómputo. La biblioteca de comunicación HCCL, co-diseñada con MTIA 300, compila cada colectivo en un conjunto de subgrafos (entradas de cola de trabajo con dependencias explícitas) que son despachados a los MEs para una ejecución autónoma, minimizando la intervención del host.
Flujo de Comunicación en MTIA 300 (Entrenamiento)
- 1 Host CPU Compila el modelo y las operaciones colectivas con HCCL.
- 2 HCCL (Compilador) Genera subgrafos de work-queue entries para colectivos.
- 3 HBM (MTIA 300) Almacena las instrucciones compiladas para los MEs.
- 4 Message Engines (MEs) Ejecutan autónomamente las operaciones colectivas.
- 5 Near-Memory Compute (NMC) Realiza reducciones a velocidad de línea (ej. AllReduce).
- 6 NIC Chiplets Manejan la comunicación RDMA dentro/entre racks.
- 7 Processing Elements (PEs) Ejecutan el cómputo de entrenamiento, aislados de la comunicación.
| Capa | Tecnología | Justificación |
|---|---|---|
| compute | MTIA 300 | Acelerador de IA diseñado específicamente para entrenamiento de modelos de ranking y recomendación, con hardware dedicado para comunicación. vs GPUs de propósito general 12x6 grid de PEs, 16 MEs dedicados, 216 GB HBM3E. |
| networking | NIC Chiplets (custom 800 Gbps RDMA) | Integrados directamente en el paquete del chip para proporcionar 1.2 TB/s de ancho de banda I/O, eliminando el cuello de botella PCIe. vs NICs externas conectadas vía PCIe 12 NICs Ethernet-based, configurables para scale-up/scale-out. |
| messaging | HCCL (Meta's Communication Library) | Biblioteca co-diseñada con el hardware para compilar y orquestar operaciones colectivas de forma autónoma en los MEs. vs NCCL (NVIDIA Collective Communications Library) Integración con PyTorch c10d y torchcomms, algoritmos topology-aware. |
| compute | Message Engines (MEs) | Unidades de procesamiento dedicadas dentro del chip para orquestar y ejecutar operaciones de comunicación, incluyendo un núcleo RISC-V y un NMC. vs Streaming Multiprocessors (GPUs) 16 MEs, cada uno con RISC-V core, NIC interface, NMC block. |
| compute | Near-Memory Compute (NMC) | Bloques de cómputo dedicados dentro de los MEs para realizar reducciones de datos a alta velocidad (128 bytes/cycle) cerca de la memoria. vs Cómputo en PEs generales Más de 2.8 TB/s de throughput de reducción colectiva. |
Trade-offs
Ganancias
- ▲ Aislamiento de cómputo y comunicación
- ▲▲ Ancho de banda de comunicación I/O
- ▲ Reducción de latencia por operación colectiva
- ▲ Eficiencia energética (implícito por menor subutilización)
Costes
- ▲ Costo de diseño y fabricación de hardware especializado
- △ Menor flexibilidad para cargas de trabajo no intensivas en comunicación
Fundamentos Teóricos
El problema de la comunicación eficiente en sistemas distribuidos a gran escala ha sido un tema central en la investigación de supercomputación y computación paralela durante décadas. Conceptos como la latencia de comunicación, el ancho de banda y la sobrecarga de sincronización son fundamentales y se abordan en trabajos clásicos sobre redes de interconexión y algoritmos de comunicación colectiva. La idea de offload de operaciones de red y cómputo de comunicación a hardware dedicado se remonta a las primeras arquitecturas de Message Passing Interface (MPI) y RDMA, donde se buscaba reducir la carga de la CPU y la latencia.
La optimización de colectivos como AllReduce es un área activa de investigación, con algoritmos como el 'doubling' o 'ring' que buscan minimizar los pasos de comunicación y maximizar el paralelismo. La integración de NICs y motores de mensajes dedicados en MTIA 300 es una aplicación directa de estos principios, llevando la optimización de la pila de comunicación al nivel del diseño de silicio, un enfoque que recuerda a los procesadores de red y las arquitecturas de 'network-on-chip' exploradas en la academia para sistemas empotrados y de alto rendimiento.