Infra behind Krea 2: How to train and serve at scale — Gabriel Jorge Menezes, Krea.ai
Krea's Gabriel Jorge Menezes shares critical insights into building the infrastructure for large-scale ML model training and serving. Key takeaways include the necessity of custom metrics beyond standard GPU utilization, aggressive checkpointing on ultra-fast storage to counter frequent cluster crashes, and a dynamic Kubernetes-based system using gang scheduling, virtual-kubelet, and taints to seamlessly shift inference workloads to external providers when training consumes on-prem GPUs. The approach highlights practical solutions for silent failures, thermal management, and optimizing resource utilization in a unified production and training environment.
Gabriel Jorge Menezes from Krea discusses the infrastructure challenges and solutions developed to train Krea 2, a diffusion transformer model trained from scratch on thousands of GPUs, and how they serve it in production. Krea 2 aims to provide creatives with tools for exploring out-of-distribution images, offering diverse styles from pixel art to photorealism, and is open-source with two checkpoints available.
Overcoming Large-Scale Training Instabilities
Training Krea 2 involved scaling from small ablation experiments to thousands of GPUs. A significant challenge encountered was the increasing frequency of crashes as the cluster size grew. These failures often manifested silently, with issues like “Nickel timeouts” where all monitoring dashboards appeared green despite communication problems. Initially, the team would meticulously debug, but they learned a pragmatic approach: “sometimes you just let it crash.” It was observed that the same nodes running the same code, after crashing, would frequently run for much longer durations (12-24 hours) on subsequent attempts. This philosophy, combined with aggressive checkpointing, was crucial for progress.
Essential Metrics for GPU Clusters
Menezes emphasizes the critical role of metrics for large-scale pre-training. Going “blind” is not an option. Key metrics and their practical implications include:
- GPU Temperature: GPUs are sensitive to heat. A single GPU running warmer than others (e.g., above 78°C) can cause throttling, instability, and erratic training behavior. The solution is to immediately remove and replace any GPU exceeding this threshold, rather than attempting to debug it.
- GPU Utilization vs. Tensor Core Utilization: Standard GPU utilization metrics are often misleading (“a lie”). They indicate the amount of time the GPU is working but not the efficiency or quality of that work. Instead, tensor core utilization is presented as a more accurate proxy for how effectively the GPU’s tensor cores are being used. It was observed that tensor core utilization directly correlated with the training image resolution, increasing as the model scaled from 128x128 to 1024x1024 pixels, indicating more demanding computations.
- InfiniBand and NVLink Metrics: These interconnect metrics are crucial but not exported by default from NVIDIA’s DCGM tools. Most communication failures in their large-scale training were cross-node (InfiniBand-related). The team built custom solutions to collect comprehensive InfiniBand metrics (throughput, message wait times, error types, packet counts). NVLink metrics, though less critical than InfiniBand, were also collected to diagnose intra-node communication errors, which could indicate issues with a specific machine.
Robust Infrastructure for Training and Serving
To manage the volatile training environment and efficiently utilize resources, Krea implemented several infrastructure strategies:
- Aggressive Checkpointing: Due to frequent crashes, continuous and fast checkpointing was essential. This required a highly performant filesystem capable of writing a terabyte of data in under 30 seconds without choking the training process. The team learned from negative experiences with Ceph, opting for a more reliable, paid solution that guarantees data integrity and speed.
- Unified Cluster with Gang Scheduling: Krea operates a single large cluster for both production inference and training.
kube-batch(Q), an open-source project, is used for gang scheduling, enabling a two-tier priority system where training workloads take precedence. Training pods have high priority and will evict inference pods if necessary. - Dynamic Inference Offloading via
virtual-kubelet: To prevent production downtime when training hogs GPUs, Krea usesvirtual-kubeletto create a “fake machine” within Kubernetes. When on-prem GPUs are fully utilized by training, inference traffic is seamlessly flipped to external providers or other clusters. This system allows for gradual migration of inference traffic, ensuring the site remains operational. TheHorizontal Pod Autoscaler(HPA) helps with resilience, allowing Kubernetes to automatically recreate failed pods, reducing manual intervention. - Smart GPU Management with Taints: Kubernetes taints are employed to manage GPU allocation dynamically. When on-prem GPUs are available and not in use by training, nodes are tainted to prevent inference pods from scheduling there, pushing traffic to external providers. When training commences and consumes all on-prem GPUs, the taint is removed, allowing new inference pods to schedule on external providers. Conversely, a descheduler gradually migrates inference pods back to the on-prem cluster when GPUs become free, avoiding a sudden rush that could destabilize production (by not using
NoExecutetaints). - Inference Robustness: Notably, diffusion transformers were found to be highly resilient in inference, capable of running on “very bad GPUs,” even those that are hot or experiencing issues.
This sophisticated infrastructure allows Krea to fully utilize its GPU cluster for critical training while maintaining stable production services, and empowers researchers by abstracting away complex resource management. Krea is actively hiring for roles to build such systems.