# Infra behind Krea 2: How to train and serve at scale — Gabriel Jorge Menezes, Krea.ai

## Executive 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: 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: 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: 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: 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: 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: 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): 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: 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]

## Practical implications

- When designing large-scale AI training infrastructure, prioritize robust monitoring of interconnects (InfiniBand/NVLink) and thermal metrics over simple GPU utilization reporting.
- Implement a layered scheduling system (like Gang scheduling) that allows high-priority research workloads to preempt lower-priority production services without causing service degradation.
- Utilize Kubernetes features like taints, tolerations, and descheduler mechanisms to dynamically manage resource allocation and prevent costly idle time or wasted compute cycles.

## Topics

AI Infrastructure, Large Language Models (LLMs), Diffusion Transformers (DiTs), Kubernetes, Distributed Computing, Krea 2, Hugging Face / GitHub

Source: https://www.youtube.com/watch?v=byn9PURoBNY