Google Cloud Blog · Barak Peleg, Asaf Ben-Tovim ·
AI21 trains its models on AI Hypercomputer
AI21 Labs utilizaba Google Cloud AI Hypercomputer con GKE para entrenar modelos de foundation como Jamba, pero enfrentaba problemas graves de fragmentación de recursos: aunque había capacidad disponible en el cluster, estaba dispersa en fragmentos demasiado pequeños para que los trabajos grandes se ejecutaran, lo que obligaba a negociaciones manuales y causaba esperas de hasta 72 horas. La adopción de Kueue como scheduler, combinada con características como Admission Fair Sharing y Topology Aware Scheduling desarrolladas en colaboración con el equipo de Kueue, automatizó completamente la asignación de recursos, reduciendo las intervenciones manuales de 20 por semana a cero y las esperas de los trabajos prioritarios de 72 horas a 12 horas. El resultado permitió que los equipos de investigación dedicaran su tiempo a experimentos e iteraciones en lugar de negociaciones de capacidad.
AI21 trains its models on AI Hypercomputer using GKE and Google Cloud A3 and A3 Ultra instances powered by NVIDIA H100 and H200 Tensor Core GPUs.