Blueprint Proposal: Distributed Training on AWS Trainium with Flyte

[samhita-alla]

Community Note

• Please vote on this issue by adding a 👍 reaction to the original issue to help the community and maintainers prioritize this request
• Please do not leave "+1" or other comments that do not add relevant new information or questions, they generate extra noise for issue followers and do not help prioritize the request
• If you are interested in working on this issue or have submitted a pull request, please leave a comment

What is the outcome that you are trying to reach?

Distributed model training can be challenging when building the entire pipeline end-to-end.

A typical setup often involves:

• Data ingestion and preprocessing for large datasets
• A distributed training execution layer
• Experiment tracking and observability
• Proper containerization for environment isolation
• Orchestration with Kubernetes (EKS)
• Durability, fault tolerance, caching, and reproducibility

While each of these components can be set up independently, stitching them together into a single, reliable, and reproducible pipeline often requires significant engineering effort. The existing examples focus narrowly on the training step, leaving orchestration, retries, caching, and multi-step workflows as a separate challenge for teams to solve.

Describe the solution you would like

I propose a fully open-source blueprint for distributed training on AWS Trainium that runs on EKS, with Flyte as the orchestration layer. For the model, I can either fine-tune BERT on an Arabic reviews dataset or pre-train BERT on the FineWeb dataset. Happy to proceed with whichever is preferred.

This blueprint would go beyond just running a training job. It would demonstrate how to manage the entire lifecycle in a reproducible and scalable way, including:

• Multi-node distributed training on Trainium
• Data streaming and preprocessing for large datasets
• Integration with PyTorch Neuron/Optimum Neuron for Trainium acceleration
• Built-in experiment tracking and observability
• Automatic caching, retries, and fault tolerance
• Integration with AWS-native services like S3 (storage), CloudWatch (logs), and ECR (images)
• Fully open-source stack: PyTorch, streaming library, Neuron SDK, Flyte, Kubeflow Training Operator, Kubernetes/EKS

By providing this orchestration-first approach, we give AI practitioners a ready-to-use blueprint that is scalable, reproducible, and easy to adapt, reducing the operational burden while keeping the infrastructure open source and AWS-native.

Describe alternatives you have considered

Additional context

[omrishiv]

Hi @samhita-alla thank you for opening the issue. This sounds like a very interesting use case that maybe we can split into multiple examples to highlight the different components of Flyte (orchestration, training, tracking, observability, etc.). We are looking for someone internally who can help review/maintain this infrastructure. Once we identify someone, we will have them comment here and we can proceed. Thank you

[samhita-alla]

Thanks @omrishiv! We’re currently working with ND Ngoka, Senior Solutions Architect at AWS, on these blueprint proposals. He’s most familiar with both Flyte and the details here.

On splitting them up: since Flyte already supports all these features natively, my view is that one comprehensive blueprint example could effectively demonstrate the features together 🙂

[omrishiv]

Hi @samhita-alla, yes, ND reached out. We're looking for someone who would be able to support the code as well and is actively using Flyte to be able to help maintain.

[omrishiv]

Hi @samhita-alla, are you able to provide an architecture diagram of what this looks like? Specifically, what we want to understand is what additional infrastructure components need to get installed into the EKS to support any further deployments. Infrastructure would be anything like operator/controllers, further deployments would be anything that would be used specifically to highlight this use case (workflows, models, things that run one off).

[samhita-alla]

hi @omrishiv, apologies for the delay. I was out of office the last couple of days. I’ll get back to you on that soon.

[samhita-alla]

@omrishiv as for the architecture diagram, we plan to use flyte 2.0 for this workflow, and since it’s relatively new compared to flyte 1.0, that’s still a work in progress. i’ll share it with you as soon as we have one ready. for reference, you can find more details on the flyte 1.0 architecture here: https://flyte.org/blog/speed-up-time-to-production-with-flyte-on-eks-from-automation-to-multicloud. does that answer your question?

since this is a distributed training workflow, we’ll be using AWS trainium as the compute backend along with the kubeflow training operator. we’re planning to upgrade to kubeflow trainer v2 or a native jobset plugin, but today, flyte 2.0 supports kubeflow training operator v1.

[omrishiv]

Hi @samhita-alla , if the diagram in the link is highlighting the components needed, yes that makes sense.