Infra behind Krea 2: How to train and serve at scale — Gabriel Jorge Menezes, Krea.ai
Summary
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]
Technical details
-
GPU Monitoring and Throttling
390s
GPUs running hotter than 78° must be immediately removed from the run to prevent throttling and instability. Focus on tensor core utilization rather than overall GPU usage percentage. [4:18]
-
Inter-Node Communication Monitoring
408s
For large-scale pre-training, custom metrics for InfiniBand and NVLink are critical because most failures stem from cross-node communication issues. These metrics must be collected manually as they are not standard exports. [6:48]
-
Cluster Resource Management (Kubernetes)
863s
The system uses Kubernetes taints and tolerations to manage GPU allocation, preventing pods from scheduling on nodes with available GPUs when training is active. This stops resource waste. [14:23]
-
Workload Scheduling
682s
A custom queue system implements Gang scheduling, giving high priority to training pods over inference pods. The architecture uses a 'virtual kubelet' concept to abstract the physical location of production traffic when GPUs are needed for research. [11:01]
Mentioned resources
- Krea 2
- Hugging Face / GitHub
Channel & topics
Watch on YouTube · Back to latest
This independent, AI-assisted summary is provided for commentary and informational purposes. It may contain errors or omit important context. Please watch the original video for the creator's complete presentation. Video, thumbnail, and related copyrights belong to their respective owners.