AWS define bloques para entrenar y servir modelos de IA

AWS publica un marco técnico para entender cómo entrenar y servir modelos fundamentales a escala usando infraestructura acelerada, redes de baja latencia y un ecosistema open source. ¿Te interesa saber dónde aparecen los cuellos de botella reales y cómo se conectan hardware, orquestación y software ML en producción? Aquí te explico lo esencial, con detalles prácticos para ingenieros y arquitectos.

Arquitectura y bloques de infraestructura

AWS organiza la pila alrededor de tres bloques acoplados: compute acelerado con mucha memoria en dispositivo, interconexión de alta banda ancha y baja latencia, y almacenamiento distribuido escalable. No es solo hardware; es la suma que permite pre-training, post-training y serving a gran escala.

• Compute acelerado: familias P5 y P6 con GPUs NVIDIA H100, H200, Blackwell B200 y B300. Los ejes de escala son rendimiento de Tensor Cores, capacidad y ancho de banda HBM, y ancho de banda de interconexión.

• Red: NVLink/NVSwitch para escala up intra-nodo y EFA (Elastic Fabric Adapter) para escala out entre nodos. EFA ofrece OS-bypass RDMA usando libfabric y el protocolo SRD, reduciendo latencia en collectives.

• Almacenamiento: jerarquía tiered con NVMe local para datos hot, FSx for Lustre para throughput paralelo y Amazon S3 para persistencia de checkpoints y datasets.

Además, AWS ofrece UltraClusters para agrupar miles de instancias con red no bloqueante y UltraServers que extienden dominios NVLink entre varias instancias, llegando a 72 GPUs y terabytes de HBM dentro de un mismo dominio NVLink. Esto cambia las reglas cuando el cuello de botella es la salida del dominio NVLink.

Orquestación y programación de recursos

Cuando un job necesita cientos o miles de GPUs, co-programarlas manualmente es imposible. Ahí entran Slurm y Kubernetes con extensiones.

• Slurm: paradigma batch clásico de HPC. Programa jobs atómicos, soporta backfill, colocación topology-aware y control fino de GPU con GRES. En AWS se despliega con ParallelCluster o con control plane gestionado como PCS y funciones de HyperPod en SageMaker.

• Kubernetes: declarativo y poderoso para despliegue, pero nativo no garantiza admission gang ni colocación topológica para collectives. Por eso hay capas adicionales:

  • Kueue: admission controller que gestiona gang scheduling y cuotas.
  • Volcano y NVIDIA KAI Scheduler: reemplazan o extienden el scheduler para placement topology-aware y gang scheduling con sensibilidad NVLink.

Amazon EKS integra device plugin NVIDIA y, junto a SageMaker HyperPod en modo EKS, añade governance, Kueue gestionado, checkpointless training, elastic training y auto-resume.

¿Qué es checkpointless training? En vez de escribir checkpoints multi-TB a almacenamiento compartido, se replica estado entre peers para recuperar fallos vía comunicación EFA. Esto reduce dependencia de IO en fallos frecuentes.

Pila software ML: desde drivers hasta frameworks avanzados

Piensa la pila como cinco capas: drivers y soporte hardware, runtimes y librerías, sustrato de comunicación, frameworks ML, y frameworks distribuidos para entrenamiento e inference.

• Drivers y runtimes: controladores NVIDIA, GDRCopy para copias CPU-GPU, driver EFA y cliente Lustre. CUDA Toolkit 13.x da soporte a arquitecturas Blackwell.

• Kernels y toolchains: optimizaciones como FlashAttention, Triton, CuTe y CUTLASS dominan el rendimiento real. Muchas ganancias vienen de kernels fusionados y especializados para atención, layernorm, MoE dispatch y KV-cache.

• Comunicación: NCCL para collectives con algoritmos topology-aware. En AWS, NCCL se conecta a libfabric vía aws-ofi-nccl para usar EFA sin cambiar la app. Para MoE, las all-to-all son críticas y pueden dominar el tiempo de paso si la paralelización de expertos escala mucho.

• Transferencias punto a punto en serving: NIXL (NVIDIA Inference Xfer Library) unifica movimientos entre HBM, DRAM y almacenamiento y se integra con backends UCX y GPUDirect Storage.

• Frameworks base: PyTorch y JAX. Esta serie se enfoca en PyTorch por su prevalencia en OSS. En PyTorch, torch.distributed ofrece process groups, DDP y FSDP2 (sharding inspirado en ZeRO).

• Frameworks de alto nivel:

  • Hugging Face Transformers + Accelerate: fácil de usar, ideal para fine-tuning y escenarios de escala moderada.
  • NVIDIA Megatron Core y NeMo: optimizados para máximo rendimiento con 3D parallelism y soporte FP8.
  • veRL: pensado para RLHF y post-training, permite mezclar backends en un mismo job.
  • vLLM y SGLang: soluciones de inference que manejan KV cache de forma paginada o con técnicas como RadixAttention para reutilizar prefijos y mejorar batching.

¿Resultado práctico? Elegir la combinación correcta entre kernels optimizados, comunicación y estrategia de paralelismo suele importar tanto o más que elegir la GPU más rápida.

Observabilidad y operación a escala

Sin telemetría no puedes operar clusters grandes. Observabilidad cubre infraestructura, workload y alertas.

• Stack estándar: Prometheus + Grafana. En AWS, AMP y AMG son las versiones gestionadas que eliminan el overhead operativo.

• Métricas clave: DCGM-Exporter para métricas GPU (utilización, HBM, ECC, XID). EFA expone counters de bytes y retransmisiones. FSx for Lustre da throughput y latencias metadata. A nivel aplicación exportas step time, tokens por segundo y pérdidas.

Un patrón común: inspeccionar DCGM para salud GPU, luego EFA/NCCL para problemas de collectives, y finalmente IO para cuellos de botella en checkpoints.

• Fallas que debes vigilar: incremento de ECC single-bit errors, aparición de XID 63 64 94 95. Estos suelen anticipar fallas graves y merecen reemplazo rápido de nodos.

• Dashboards y alertas: el dashboard GPU Health - Cluster es un buen punto de partida. Parte 5 de la serie profundiza en reglas de alerta y retención de métricas.

Qué significa todo esto para tu proyecto

Si trabajas con modelos grandes, la lección es clara: los tres regímenes de escala - pre-training, post-training y test-time - convergen en los mismos requisitos infraestructurales. No se trata solo de comprar GPUs; se trata de alinear dominio NVLink, EFA, almacenamiento y orquestación con las necesidades de comunicación de tu modelo.

Preguntas prácticas que debes hacerte:

• Tu paralelismo exige all-to-all intensas o puedes quedarte en data/tensor parallelism?
• Necesitas checkpointless recovery o tu pipeline tolera IOPS altos para checkpoints multi-TB?
• Tu scheduler respeta colocación topology-aware para minimizar hops de red?

Entender estos puntos de integración es la base para diagnosticar cuellos de botella y tomar decisiones de escalado inteligentes.

Fuente original

https://huggingface.co/blog/amazon/foundation-model-building-blocks