Infra behind Krea 2: How to train and serve at scale — Gabriel Jorge Menezes, Krea.ai
Gabriel Jorge Menezes details the complex infrastructure required to train and serve Krea 2, a diffusion transformer model trained from scratch on thousands of GPUs. The system addresses challenges like silent failures at scale, GPU thermal throttling, and cross-node communication issues by implementing advanced monitoring (tensor core utilization, InfiniBand metrics). For serving, they built a robust architecture using Gang scheduling and Kubernetes features (virtual kubelet, taints/tolerations) to ensure training workloads can utilize the entire cluster while maintaining production uptime through seamless traffic flipping.
Key takeaways
-
Metrics are essential for large-scale pre-training
9:50
Do not rely on GPU utilization (which is 'a lie'). Instead, monitor tensor core utilization and collect custom metrics like InfiniBand/NVLink errors, as most failures relate to cross-node communication. [5:58], [6:48]
-
Embrace failure for stability
4:18
When scaling training runs, instead of debugging every crash, it is often more efficient to 'let it crash.' The system should be designed to recover and run successfully on the same nodes over extended periods. [4:18]
-
Checkpointing must be extremely fast
8:29
To make long training runs survivable, checkpoint aggressively against a high-speed filesystem capable of writing terabytes quickly (e.g., achieving >1 TB/30 seconds). [8:29]
-
Decouple training and production workloads
11:01
Use a system that allows high-priority training jobs to utilize the entire cluster while seamlessly migrating inference traffic (production) to external providers or other clusters, ensuring zero downtime. [11:01]