Scaling Custom LLMs on EKS: Trainium and Inferentia2-Powered AI Infrastructure

I attended AWS Community Day Kochi on 20th December 2025 and it was full of amazing insights. There were a lot of great seminars throughout the day on different AWS services but one session in particular totally caught my eye. The session featured a technical deep dive, “Scaling Custom LLMs on EKS: Trainium and Inferentia2-Powered AI Infrastructure”.

Here’s the rundown of what was covered and how it all works and why it matters if you’re building AI apps today.

The Challenge: When Do You Outgrow Standard APIs?

Let's face it, using off-the-shelf APIs like OpenAI or Anthropic is great for rapid prototyping. You obtain an API key, you write three lines of code, and you have intelligence embedded right into your program. But what if your app goes viral and takes off?

The speaker hit us with a hard truth straight away: API costs increase dramatically with request volume. Think of it as a cab. It is quite useful for short, occasional journeys. However, if you are driving 100 miles each day, then owning your own automobile makes a lot more financial sense.

Session data suggests the important inflection point is at around 10 million requests each day. That’s when third party APIs become too expensive and building a custom platform is a very practical, affordable necessity.

Enter AWS Purpose-Built Chips

This is a huge scalability problem. To overcome this AWS has come up with purpose-made hardware accelerators that are specifically developed and constructed from the ground up for Generative AI workloads.

We all know what a typical GPU is, but they are expensive to rent and hard to get at scale. AWS has created its own silicon to solve this very problem.

• AWS Inferentia (Inf1): It offers the lowest cost per inference in the cloud for Deep Learning models, with costs up to 70% cheaper than normal EC2 instances.

• AWS Inferentia2 (Inf2): Specifically designed for large language model (LLM) and diffusion model. It offers up to 40% better price performance than equivalent EC2 instances.

• AWS Trainium (Trn1): Your powerhouse for model training. This means you can achieve up to 50% savings on training costs compared to equivalent Amazon EC2 instances.

The Architecture: Bringing it All Together on EKS

Great software is essential to make hardware useful. This is where the AWS Generative AI stack and SDK for Neurone come into play. Neurone SDK bridges third-party applications such as PyTorch, Ray and vLLM to the underlying Trainium and Inferentia hardware.

The worldwide plan showed a very resilient arrangement across regions. Users connect to an Application Load Balancer via a Global Accelerator, and the Application Load Balancer sends the traffic to a Karpenter-managed, dynamically provisioned Amazon EKS cluster. The models are directly served from the hardware nodes using an Nvidia Triton Inference Server in the cluster, running vLLM and the Neurone ML SDK.

The Reality of Training and Serving at Scale

Imagine painting a big canvas all by yourself with a teeny little brush. That's about how it feels training a massive LLM on a single GPU.

The speaker states that training a certain large model on a single GPU may take ~120 days and cost $500k+. A Trainium Cluster provides the ability to run numerous processors in simultaneously. That takes overall training time down to just 14 days, and cuts the cost to $150k - an astounding 70% cheaper.

Just as vital is serving. 500ms latency is a sluggish experience for traditional GPUs. Inferentia2 lowers this down to a 100-200ms latency, making the experience feel nearly quick, all while being 70% cheaper.

Hands-On: The EKS Configurations and Code

This is where it gets good. The speaker dived straight into the live demo, presenting the real setups that make this magic happen.

1. The Kubernetes Magic Line

If you want to use a Trainium or Inferentia device in your Kubernetes deployment, just add one magic line in your resource limits: [aws.amazon.com/neuron](https://aws.amazon.com/neuron): "1" .

Here is what the manifest looks like:

# The Magic Line - Requesting Neuron Device
resources:
  requests:
    cpu: "4"
    memory: 16Gi
    aws.amazon.com/neuron: "1"   # <- Request 1 Trainium/Inferentia device
  limits:
    cpu: "8"
    memory: 32Gi
    aws.amazon.com/neuron: "1"

To bridge the gap between Kubernetes and the physical hardware, you deploy a DaemonSet called the neuron-device-plugin. It automatically discovers the hardware and mounts the /dev/neuron* paths:

# DaemonSet discovers and exposes Neuron devices
apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: neuron-device-plugin
  namespace: kube-system
spec:
  template:
    spec:
      containers:
        - name: neuron-device-plugin
          image: public.ecr.aws/neuron/neuron-device-plugin:2.19.16.0
          securityContext:
            privileged: true  # Required to access /dev/neuron*
          volumeMounts:
            - name: device-plugin
              mountPath: /var/lib/kubelet/device-plugins
            - name: neuron-dir
              mountPath: /dev  # Access Neuron devices

2. Protecting Your Budget with Taints

Here’s a huge money saving tip. Trainium on-demand nodes cost around $1.34 per hour. You do not want to run simple system pods like CoreDNS on those pricey ML instances!

Using Terraform, you can apply a taint to block non-ML pods:

# Terraform: Node Group with Taint
taint {
  key    = "aws.amazon.com/neuron"
  value  = "true"
  effect = "NO_SCHEDULE" # Block non-ML pods!
}

Then you just add a toleration in your kubernetes deployment so your ml pods can schedule there:

# Kubernetes: Pod must tolerate to run on Neuron node
tolerations:
  - key: "aws.amazon.com/neuron"
    operator: "Equal"
    value: "true"
    effect: "NoSchedule"

nodeSelector:
  node-type: trainium          # Target Trainium nodes
  aws.amazon.com/neuron: "true"

3. Training with PyTorch + Neuron

The real training code with Pytorch and Hugging Face is delightfully simple. Notice how loading the model in bfloat16 offers you a 2x memory save right away:

# Training on Trainium
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load model
model = AutoModelForCausalLM.from_pretrained(
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    torch_dtype=torch.bfloat16,  # BF16 = 2x memory savings
)

# Assume model_path is defined
tokenizer = AutoTokenizer.from_pretrained(model_path)
tokenizer.pad_token = tokenizer.eos_token

# Simple training loop
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
model.train()

for step, batch in enumerate(dataloader):
    outputs = model(
        input_ids=batch["input_ids"],
        attention_mask=batch["attention_mask"],
        labels=batch["labels"],
    )
    loss = outputs.loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

    print(f"Step {step} - Loss: {loss.item():.4f}")

4. Serving the Model with FastAPI

The speaker exhibited a clean FastAPI wrapper providing the trained model on Inferentia2.

# Inference Server on Inferentia2
from fastapi import FastAPI
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

app = FastAPI(title="TinyLlama Inference API")

# Load model once at startup
model = AutoModelForCausalLM.from_pretrained("/models/tinyllama")
tokenizer = AutoTokenizer.from_pretrained("/models/tinyllama")

@app.get("/health")
async def health():
    return {"status": "healthy"}

@app.post("/v1/generate")
async def generate(prompt: str, max_tokens: int = 100):
    inputs = tokenizer(prompt, return_tensors="pt")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=0.7,
            do_sample=True,
        )

    return {
        "generated_text": tokenizer.decode(outputs[0], skip_special_tokens=True)
    }

The Bottom Line: Math and Money

Your workload is the single most important factor for your memory needs. For a 7B parameter model running in BF16:

• Inference Memory: Needs about 16.8 GB of memory. This fits nicely on an inf2.xlarge instance.

• Training Memory: You need a big 56 GB because you need to save gradients and optimisers.

The final pricing comparison was the most jaw-dropping slide. APIs with significant traffic can easily run to $60,000/month off-the-shelf. Running your own EKS cluster with Inf2 On-Demand instances reduces that to $547.

But the true game-changer is that you are able to set up Spot Instances with Terraform.

# Terraform: Enable Spot Instances
resource "aws_eks_node_group" "trainium" {
  instance_types = ["trn1.2xlarge"]

  # THE MONEY SAVER 
  capacity_type = "SPOT"  # vs "ON_DEMAND"

  # Scale to zero when not training
  scaling_config {
    desired_size = 0      # Start at zero!
    min_size     = 0      # Allow scale to zero
    max_size     = 1
  }
}

EKS with Inf2 Spot instances (which can scale down to zero when idle) brings the monthly cost down to a mere $166. That's a crazy inexpensive $0.17 every 1K inferences!

Key Takeaways

• The 10M Threshold: Third party LLM APIs stop making economical sense and specialised infrastructure becomes needed after 10 million daily requests.

• Massive Cost Slashes: Move to an EKS cluster with Trainium and Inferentia2 and cut 50% of your training costs and 40% of your inference costs.

• Budget Guardrails: Kubernetes standards and tolerations prohibit basic system pods from schedule-squatting on your premium $1.34/hr ML nodes.

• Performance Without Compromise: Lowering infrastructure expenses doesn’t mean losing user experience. Inferentia2 readily maintains a latency of 100-200ms and a sub-100ms P99 latency even at a large 70B parameter scale.

Conclusion

Building custom Generative AI infrastructure is no longer some black art for the tech giants with endless cash at the end of the day. AWS has democratised the hardware layer with purpose-built silicon, and the software gap is bridged with open-source tools like the Neurone SDK, making it developer-friendly.

Long story short, if you are establishing an AI firm or trying to extend an engineering staff beyond basic API wrappers, it's well worth the effort to break free from third-party locks by understanding this purpose-built AWS stack.

About the Author

As an AWS Community Builder, I enjoy sharing the things I've learned through my own experiences and events, and I like to help others on their path. If you found this helpful or have any questions, don't hesitate to get in touch! 🚀

🔗 Connect with me on LinkedIn

References

Event: AWS Community Day Kochi

Topic: Scaling Custom LLMs on EKS: Trainium and Inferentia2-Powered AI Infrastructure

Date: December 20, 2025

Also Published On

AWS Builder Center

Hashnode