Compiling 70B+ LLMs with NxDI on AWS Trainium using EKS

A customer aiming to serve models with tens of billions of parameters—while saving up to 50% on accelerated compute costs—can use EC2 Trainium alongside EKS for scalable, efficient container orchestration and seamless AWS integration, and leverage real-time monitoring and proactive troubleshooting with services like CloudWatch or OpenTelemetry. Model compilation is memory intensive and requires optimal configuration to maximize resources while ensuring sufficient capacity for monitoring tools. This post demonstrates how to compile the DeepSeek R1 70B distilled model on a trn1.32xlarge EC2 instance with optimal tensor parallelism and NxDI compilation, while monitoring progress using CloudWatch ContainerInsights. We recommend reviewing "Get started with DeepSeek R1 on AWS Inferentia and Trainium" before reading this guide. For more details on NxDI, please refer to the NxDI Overview.

Also configure EC2NodeClass that uses the EKS-managed DLAMI and allocate large disks-space to store the compiled model artifacts (graph and weights)

Note that setting resources.limits.memory to "465Gi" is crucial because the default memory requests are much lower than what the compilation process requires. In Kubernetes, the admission controller and kubelet eviction logic monitor memory usage closely. If a pod exceeds its allocated memory or if node-level memory pressure is detected, the kubelet’s eviction manager may terminate the pod to reclaim resources. This high limit ensures that the model compilation process—which involves memory-intensive operations—is not preempted by other processes (like monitoring applications) competing for resources, thereby preventing unwanted evictions by the control plane.

The number of spikes in the memory utilization expected be as the value of the configured TENSOR_PARALLEL_SIZE
By the end pf this job expects to see the model in your HuggingFace user.

You can also track the progress via kubectl: