Distributed training of large language models on AWS Trainium

By Xinwei Fu, Zhen Zhang, Haozheng Fan, Guangtai Huang, Randy Huang, Rahul Solanki, Fei Wu, Ron Diamant, Yida Wang
2024
Download Copy BibTeX
Copy BibTeX
Large language models (LLMs) are ubiquitously powerful but prohibitively expensive to train, often requiring thousands of compute devices, typically GPUs. To reduce the cost of training LLMs for customers, Amazon Web Services (AWS) launched the Amazon EC2 trn1 instances, powered by AWS Trainium, Amazon’s homegrown deep-learning accelerator, as an alternative to distributed LLM training. The trn1 instances provide a high-performance LLM training solution at a lower cost compared to their GPU-based counterpart, the p4d instances, which are powered by Nvidia A100 GPUs. This paper describes the design and development of the Neuron Distributed Training Library, a component of the AWS Neuron SDK, which enables distributed training of large language models on AWS Trainium. Neuron Distributed Training Library supports a variety of existing distributed training techniques with unified interfaces, and provides features to address trn1-specific challenges as well. Our evaluation shows that trn1 instances, specifically the trn1.32xlarge, achieve better or comparable performance (up to 24.6% improvement) while offering significant lower costs (up to 46.3% cost saving) in selected workloads when compared to p4d.24xlarge instances. As a result, AWS Trainium has been adopted for training numerous external and internal models, showcasing its high-performance and cost effectiveness. Several supported open-source LLMs are accessible via HuggingFace Optimum Neuron.
Research areas

Latest news

IL, Tel Aviv
Are you a scientist interested in pushing the state of the art in Information Retrieval, Large Language Models and Recommendation Systems? Are you interested in innovating on behalf of millions of customers, helping them accomplish their every day goals? Do you wish you had access to large datasets and tremendous computational resources? Do you want to join a team of capable scientist and engineers, building the future of e-commerce? Answer yes to any of these questions, and you will be a great fit for our team at Amazon. Our team is part of Amazon’s Personalization organization, a high-performing group that leverages Amazon’s expertise in machine learning, generative AI, large-scale data systems, and user experience design to deliver the best shopping experiences for our customers. Our team is building next-generation personalization systems powered by Large Language Models. We are tackling novel research challenges to help customers discover products they'll love - at Amazon scale and latency requirements. We are a team uniquely placed within Amazon, to have a direct window of opportunity to influence how customers will think about their shopping journey in the future. As an Applied Science Manager, you will lead a team of scientists working at the frontier of LLM-based personalization. You will set the technical vision, drive the research agenda, and ensure your team delivers production-ready solutions. You will hire, mentor, and develop world-class scientists while fostering a culture of innovation and scientific rigor. You will partner closely with engineering and product teams to translate ambitious research into customer-facing impact, and represent your team's work to senior leadership. Please visit https://www.amazon.science for more information.
BR, SP, Sao Paulo
Do you feel the challenge and the adrenaline kick when a huge data-set stares you in the face and you know that somewhere inside are hidden very important business insights that can fundamentally alter the way top business leaders think and act? Do you enjoy presenting strong data backed insights to business leaders; insights that can topple their long held beliefs and compel them to change their direction completely? If yes, then you are the one we are looking for. We are looking to invite passionate leaders, with expertise in generate power business insights from very large datasets, on a journey where the primary aim would be to enable needle moving business impacts through statistical analysis. We are looking for leaders who can envision the design and development of analytical infrastructure which can support strategic and tactical decision-making. Those who join this high visibility team would have to navigate through significant ambiguity in defining business problems and converting them to analytical problems. This role requires additional exposure and experience to Machine Learning. Key job responsibilities Use machine learning and analytical techniques to create scalable solutions for business problems • Analyze and extract relevant information from large amounts of Amazon’s historical business data to help automate and optimize key processes • Design, development, evaluate and deploy innovative and highly scalable ml models such as risk scorecards, income models, fraud models for predictive learning in credit risk applications • Research and implement novel machine learning and statistical approaches • Work closely with software engineering teams to drive real-time model implementations and new feature creations • Work closely with business owners and operations staff to optimize various business operations • Establish scalable, efficient, automated processes for large scale data analyses, model development, model validation and model implementation • Mentor other scientists and engineers in the use of ML techniques • Innovate with the latest GenAI technology to build highly automated solutions for efficient customer promotions • Design, develop and deploy end-to-end machine learning solutions in the Amazon production environment to delight Amazon customers • Collaborate with cross-functional teams to develop comprehensive ML/statistical models that can scale to millions of customers to multiple countries Understand the credit risk data and evaluate the best ml model/ solution for dynamic business problems. About the team Brazil Payments is part of the International Emerging Stores Payments team and focuses on supporting the launch of new payment and financial products to our customers in Brazil.
US, NY, New York
Amazon Advertising drives billions of ad impressions and millions of clicks daily, powering discovery and sales for advertisers across Amazon's Retail and Marketplace businesses. The Ads Marketing Decision Science team sits at the intersection of data science and marketing strategy. We build intelligent, data-driven systems that analyze advertiser behavior at large scale to deliver the right guidance to the right advertiser at the right time. Our work spans behavioral modeling, content intelligence, automated decision systems, and GenAI applications, enabling personalized marketing experiences that help advertisers make smarter advertising decisions and grow their business on Amazon. We are looking for a Data Scientist who brings strong fundamentals in machine learning, causal inference, and statistical modeling to solve real advertiser problems. You will build predictive models, design experiments, develop segmentation frameworks, and leverage GenAI capabilities where applicable, taking solutions end-to-end from proof-of-concept to production at scale. You will partner closely with scientists, engineers, and product managers on a daily basis to prototype rapidly, ensure data integrity in production systems, and deliver measurable advertiser impact. If you are passionate about solving real-world problems with next level science, come join us as we innovate and make history. Key job responsibilities • Define and execute data science solutions end-to-end, from problem framing through production deployment. • Build machine learning models (classification, regression, clustering, ranking) for advertiser segmentation, propensity modeling, and recommendations. • Apply causal inference and experimentation methods (A/B testing, difference-in-differences, propensity score matching) to measure the impact of marketing interventions. • Analyze large-scale advertiser behavioral data to identify trends, surface growth opportunities, and support optimal decision making. • Collaborate with colleagues across science and engineering disciplines for fast turnaround proof-of-concept prototyping at scale. • Establish and drive data hygiene best practices to ensure coherence and integrity of data feeding into production ML/AI solutions. • Leverage GenAI and LLM capabilities to enhance science products where applicable A day in the life You will solve real-world problems by analyzing large volumes of advertiser data, building predictive models, designing experiments, and measuring business impact. You will prototype rapidly, validate ideas with data, and partner with engineers to productize and scale successful solutions. You will collaborate daily with scientists, engineers, and product managers across the advertising organization, working in a cross-functional, fast-paced environment where data drives decisions and helps advertisers grow. About the team We are a team of Applied Scientists, Research Scientists, Data Scientists, and Business Intelligence Engineers with deep expertise in ML, NLP, Gen-AI, RL, and causal inference, from a diverse range of backgrounds. We partner closely with strong engineers, product managers, and sales leaders who bring ads-industry depth and experience building scalable modeling and software solutions.
ES, Madrid
We're building the intelligence behind how customers discover, trust and enjoy products on Amazon. We're solving complex catalogue quality challenges with machine learning, enhancing product discovery through computer vision and multimodal AI and pioneering agentic systems that autonomously navigate and stress-test the Amazon shopping experience to surface insights at scale. We're looking for PhD students across multiple research domains to invent, design, and implement state of the art solutions for never before solved problems. Your work here won't just stay in a notebook, it ships to production and reaches customers worldwide. Check out the details below including the job responsibilities, team details, and basic qualifications before submitting your application. You can find more information about the Amazon Science community as well as interview preparation tips via the links below; - https://www.amazon.science/ - https://amazon.jobs/content/en/career-programs/university/science - https://amazon.jobs/content/en/how-we-hire/university-roles/applied-science Key job responsibilities As an Applied Science Intern, you will own the design and development of end-to-end systems. You'll have the opportunity to write technical white papers, create roadmaps and drive production level projects that will support Amazon Science. You will work closely with Amazon scientists and other science interns to develop solutions and deploy them into production. You will have the opportunity to design new algorithms, models, or other technical solutions whilst experiencing Amazon's customer focused culture. The ideal intern should have the ability to work with diverse groups of people and cross-functional teams to solve complex business problems. A day in the life You'll spend your first weeks scoping your project with your mentor, then own the research and implementation end-to-end. Your work could involve developing machine learning and data analysis solutions that detect and resolve catalogue quality issues at massive scale, building computer vision and multimodal learning models that transform how customers discover and interact with products, engineering universal ML systems that make shopping on Amazon easier and more visually delightful, or creating autonomous agentic shoppers that tirelessly navigate the Amazon website to provide feedback and actionable insights — depending on the team you're matched with. Many interns publish at top-tier conferences or see their work deployed to production before the internship ends. Interns may also be considered for a return offer at the end of their internship, subject to performance evaluation and headcount availability. Further benefits of an Amazon Science internship include; - All of our internships offer a competitive salary - Interns are paired with an experienced manager and mentor(s) - Interns get invited to different intern program or office events - Interns can build their professional and personal network with other Amazon Scientists - Interns can potentially publish work at top tier conferences About the team We're hiring interns for multiple teams in Spain including but not limited to; •Tamale - Using Machine Learning and Data analysis solutions to solve complex catalogue quality problems • NintAI- Developing AI solutions, focusing on computer vision and multimodal learning to enhance how customers discover and interact with products • Home Innovation tech- Building universal, state of the art Machine Learning technology that makes shopping on Amazon easier and more visually delightful for our customers • EU Intech- Pioneers a population of agentic shoppers, autonomous AI agents, that tirelessly navigate and shop on the Amazon website, providing feedback and insights to improve the customer experience You'll submit a single application and we'll match you with science teams best aligned with your research interests. Applications are reviewed on a rolling basis, and your application stays active until we find a team match or confirm there are no matches available. Start dates are available throughout the year for durations of between 3–6 months. Please note, each team has different start date and duration preferences — your recruiter will confirm the preferences of the team you're matched with prior to interviewing. We offer science internships in multiple locations across the EMEA region and you can indicate your interest in all these locations by applying here (Austria, Estonia, France, Germany, Ireland, Israel, Italy, Jordan, Luxembourg, Netherlands, Poland, Romania, South Africa, Spain, Sweden, UAE, and UK). Please note we do not offer remote internships.
GB, London
We're building systems that turn research into real-world impact — ensuring flawless streaming for millions of Prime Video customers, shaping the future of AI-driven shopping with Rufus, advancing speech generation and conversational AI, optimizing large-scale infrastructure through intelligent observability, and transforming how people and jobs find each other. If this sounds interesting to you then we have a range of opportunities for you to explore. We're looking for PhD students across multiple research domains to invent, design, and implement state-of-the-art solutions for never-before-solved problems. Your work here won't just stay in a notebook — it ships to production and reaches customers worldwide. Check out the details below including the job responsibilities, team details, and basic qualifications before submitting your application. You can find more information about the Amazon Science community as well as interview preparation tips via the links below; - https://www.amazon.science/ - https://amazon.jobs/content/en/career-programs/university/science - https://amazon.jobs/content/en/how-we-hire/university-roles/applied-science Key job responsibilities As an Applied Science Intern, you will own the design and development of end-to-end systems. You'll have the opportunity to write technical white papers, create roadmaps and drive production level projects that will support Amazon Science. You will work closely with Amazon scientists and other science interns to develop solutions and deploy them into production. You will have the opportunity to design new algorithms, models, or other technical solutions whilst experiencing Amazon's customer focused culture. The ideal intern should have the ability to work with diverse groups of people and cross-functional teams to solve complex business problems. A day in the life You'll spend your first weeks scoping your project with your mentor, then own the research and implementation end-to-end. Your work could involve building computer vision models that detect quality issues across Prime Video's content library, developing multimodal AI solutions that power Amazon's shopping assistant Rufus, creating automated reasoning tools that verify distributed systems at scale, engineering intelligent observability frameworks for large-scale infrastructure, or advancing recommendation models that connect people with the right jobs — depending on the team you're matched with. Many interns publish at top-tier conferences or see their work deployed to production before the internship ends. Interns may also be considered for a return offer at the end of their internship, subject to performance evaluation and headcount availability. Further benefits of an Amazon Science internship include; - All of our internships offer a competitive salary - Interns are paired with an experienced manager and mentor(s) - Interns get invited to different intern program or office events - Interns can build their professional and personal network with other Amazon Scientists - Interns can potentially publish work at top tier conferences About the team We're hiring interns for multiple teams in the UK, including but not limited to; • Prime Video Video Quality Analysis – Develops AI and machine learning solutions using computer vision, audio processing, and generative AI to detect and prevent streaming quality issues across Prime Video's vast content library • Rufus Features Science UK – Shapes AI-driven shopping experiences at Amazon, working on projects from enabling Rufus to take actions on behalf of customers to generating multimodal answers combining text, image, audio, and video. • READI – Observability, Triage & Peak Readiness – Builds intelligent log analytics and automated performance frameworks that transform system telemetry into actionable insights for large-scale distributed systems. • Automated Reasoning Group – Ensures program and systems correctness through deductive proof, model checking, formal verification, and runtime conformance monitoring for large-scale distributed systems. You'll submit a single application and we'll match you with science teams best aligned with your research interests. Applications are reviewed on a rolling basis, and your application stays active until we find a team match or confirm there are no matches available. Start dates are available throughout the year for durations of between 3–6 months. Please note, each team has different start date and duration preferences — your recruiter will confirm the preferences of the team you're matched with prior to interviewing. We offer science internships in multiple locations across the EMEA region and you can indicate your interest in all these locations by applying here (Austria, Estonia, France, Germany, Ireland, Israel, Italy, Jordan, Luxembourg, Netherlands, Poland, Romania, South Africa, Spain, Sweden, UAE, and UK). Please note we do not offer remote internships.
IL, Tel Aviv
We're building systems that turn research into real-world impact — powering sports experiences for millions of Prime Video customers, pioneering multimodal document intelligence, building autonomous AI agents that reason, plan, and act, and transforming how customers discover products they love. If this sounds interesting to you then we have a range of opportunities for you to explore. We're looking for PhD students across multiple research domains to invent, design, and implement state-of-the-art solutions for never-before-solved problems. Your work here won't just stay in a notebook — it ships to production and reaches customers worldwide. Check out the details below including the job responsibilities, team details, and basic qualifications before submitting your application. You can find more information about the Amazon Science community as well as interview preparation tips via the links below; - https://www.amazon.science/ - https://amazon.jobs/content/en/career-programs/university/science - https://amazon.jobs/content/en/how-we-hire/university-roles/applied-science Key job responsibilities As an Applied Science Intern, you will own the design and development of end-to-end systems. You'll have the opportunity to write technical white papers, create roadmaps and drive production level projects that will support Amazon Science. You will work closely with Amazon scientists and other science interns to develop solutions and deploy them into production. You will have the opportunity to design new algorithms, models, or other technical solutions whilst experiencing Amazon's customer focused culture. The ideal intern should have the ability to work with diverse groups of people and cross-functional teams to solve complex business problems. A day in the life You'll spend your first weeks scoping your project with your mentor, then own the research and implementation end-to-end. Your work could involve building computer vision models that power live sports experiences for Prime Video, developing multimodal GenAI solutions for AWS document intelligence, creating agentic AI systems that reason and act autonomously, or advancing recommendation models that transform how customers discover content — depending on the team you're matched with. Many interns publish at top-tier conferences or see their work deployed to production before the internship ends. Interns may also be considered for a return offer at the end of their internship, subject to performance evaluation and headcount availability. Further benefits of an Amazon Science internship include; - All of our internships offer a competitive salary - Interns are paired with an experienced manager and mentor(s) - Interns get invited to different intern program or office events - Interns can build their professional and personal network with other Amazon Scientists - Interns can potentially publish work at top tier conferences About the team We're hiring interns for multiple teams in Israel, including but not limited to; • Prime Video Sports — Build innovative sports experiences for Prime Video, spanning computer vision, 3D simulation, and personalized content recommendations. • Personalization — Leverage LLMs, NLP, and recommender systems to match customers with products that align with their passions and shopping preferences. • Agentic AI — Build next-generation agentic AI systems that automate real-world knowledge work, spanning retrieval-augmented generation, multi-agent workflows, long-term memory and personalization, knowledge-graph construction, and agent evaluation. • DS3 Textract — Develop multimodal generative AI algorithms that pioneer state-of-the-art document understanding solutions impacting millions of customers. You'll submit a single application and we'll match you with science teams best aligned with your research interests. Applications are reviewed on a rolling basis, and your application stays active until we find a team match or confirm there are no matches available. Start dates are available throughout the year. Some teams offer full-time internships (3–6 months) while others offer part-time positions (50–60%, 8–12 months) — your recruiter will confirm the format during team matching. We offer science internships in multiple locations across the EMEA region and you can indicate your interest in all these locations by applying here (Austria, Estonia, France, Germany, Ireland, Israel, Italy, Jordan, Luxembourg, Netherlands, Poland, Romania, South Africa, Spain, Sweden, UAE, and UK). Please note we do not offer remote internships.
US, NY, New York
Our team, AWS Central Econ and Science, partners across AWS leveraging data, economics, science, and business context to build new tech, tools and policies that help AWS customers and the business. In this role you will help build out AWS's CLV framework in support of direct, partner and marketplace selling motions. This includes end to end ownership of notification for human and agentic selling motions to improve customer value from AWS products. Key job responsibilities Ownership includes developing closed loop causal measurement systems along with explainability layers. We work backward through partnering deeply with our business partners to improve processes in addition to science products. We want this role to shape the direction of our efforts to drive value of AWS customers. A day in the life Our team takes big swings and works on hard cross organizational problems where the optimal success rate is not 100%. We ask people to grow their skills and stretch and do so in a supportive and fun environment. It’s about empirically measured impact, advancement, and fun on our team. We work hard during work hours but we also don’t encourage working at nights or on weekends except in very rare, high stakes cases. Burn out isn’t a successful long run strategy. Because we invest in the long run success of our group it’s important to have hobbies, relax and then come to work refreshed and excited. It makes for bigger impact, faster skill accrual and thus career advancement. About the team Our group is technically rigorous and encourages ongoing academic conference participation and publication. Our leaders are here for you and to enable you to be successful. We believe in being servant leaders focused on influence: good data work has little value if it doesn’t translate into actionable insights that are rolled out and impact the real economy. We are communication centric since being able to explain what we do ensures high success rates and lowers administrative churn. Also: we laugh a lot. If it’s not fun, what’s the point?
US, NY, New York
We are seeking an Applied Scientist to lead the development of evaluation frameworks and data collection protocols for robotic capabilities. In this role, you will focus on designing how we measure, stress-test, and improve robot behavior across a wide range of real-world tasks. Your work will play a critical role in shaping how policies are validated and how high-quality datasets are generated to accelerate system performance. You will operate at the intersection of robotics, machine learning, and human-in-the-loop systems, building the infrastructure and methodologies that connect teleoperation, evaluation, and learning. This includes developing evaluation policies, defining task structures, and contributing to operator-facing interfaces that enable scalable and reliable data collection. The ideal candidate is highly experimental, systems-oriented, and comfortable working across software, robotics, and data pipelines, with a strong focus on turning ambiguous capability goals into measurable and actionable evaluation systems. Key job responsibilities - Design and implement evaluation frameworks to measure robot capabilities across structured tasks, edge cases, and real-world scenarios - Develop task definitions, success criteria, and benchmarking methodologies that enable consistent and reproducible evaluation of policies - Create and refine data collection protocols that generate high-quality, task-relevant datasets aligned with model development needs - Build and iterate on teleoperation workflows and operator interfaces to support efficient, reliable, and scalable data collection - Analyze evaluation results and collected data to identify performance gaps, failure modes, and opportunities for targeted data collection - Collaborate with engineering teams to integrate evaluation tooling, logging systems, and data pipelines into the broader robotics stack - Stay current with advances in robotics, evaluation methodologies, and human-in-the-loop learning to continuously improve internal approaches - Lead technical projects from conception through production deployment - Mentor junior scientists and engineers
US, WA, Seattle
Prime Video is a first-stop entertainment destination offering customers a vast collection of premium programming in one app available across thousands of devices. Prime members can customize their viewing experience and find their favorite movies, series, documentaries, and live sports – including Amazon MGM Studios-produced series and movies; licensed fan favorites; and programming from Prime Video add-on subscriptions such as Apple TV+, Max, Crunchyroll and MGM+. All customers, regardless of whether they have a Prime membership or not, can rent or buy titles via the Prime Video Store, and can enjoy even more content for free with ads. Are you interested in shaping the future of entertainment? Prime Video's technology teams are creating best-in-class digital video experience. As a Prime Video technologist, you’ll have end-to-end ownership of the product, user experience, design, and technology required to deliver state-of-the-art experiences for our customers. You’ll get to work on projects that are fast-paced, challenging, and varied. You’ll also be able to experiment with new possibilities, take risks, and collaborate with remarkable people. We’ll look for you to bring your diverse perspectives, ideas, and skill-sets to make Prime Video even better for our customers. With global opportunities for talented technologists, you can decide where a career Prime Video Tech takes you! The Prime Video Title Lifecycle Presentation team sits at the intersection of science, experimentation, and customer experience. We leverage data signals and rigorous testing to present the most engaging information about our content to customers at precisely the right moment. Our mission is to ensure every customer interaction with Prime Video content is informed, relevant, and compelling in order to drive discovery and engagement across our vast catalog. We're seeking a Sr. Applied Scientist who excels at building sophisticated machine learning systems for content presentation and discovery. The ideal candidate brings deep expertise in: - Multi-modal embeddings for rich metadata representation, enabling nuanced understanding of content attributes and customer preferences - Contextualized ranking systems that adapt to customer intent, viewing context, and real-time signals - Reinforcement learning frameworks that create continuous improvement loops, allowing our systems to learn and optimize from customer interactions over time - General modeling techniques with strong fundamentals in machine learning and statistical methods - Recommender systems experience, with proven ability to build and scale personalization solutions You'll work with technology to solve complex problems in content discovery, leveraging large-scale data to create experiences that delight millions of Prime Video customers worldwide. Key job responsibilities - Lead Cross-Functional Science Initiatives: Drive a diverse portfolio of applied science projects spanning recommender systems, generative AI agent development and evaluation across multiple modalities, and computer vision applications. Demonstrate both breadth of understanding across technical domains and sufficient depth in each area to effectively lead multiple concurrent initiatives to successful outcomes. - Bridge Science and Engineering for Production-Scale Deployment: Partner with engineering teams to productionize machine learning models at Prime Video scale. Develop production-ready science code that meets engineering standards for performance, reliability, and maintainability, ensuring seamless transition from research to deployment. - Mentor and Develop Technical Talent: Provide technical mentorship and guidance to junior scientists and engineers on applied science methodologies, best practices, and professional development. Foster a culture of scientific rigor and continuous learning within the team.
US, CA, Sunnyvale
We are seeking an Applied Scientist to focus on Robot Navigation. In this role, you'll research and develop advanced navigation systems that enable robots to move reliably and safely through complex, dynamic environments. You'll work across a broad spectrum of navigation approaches—from classical methods to learning-based techniques and foundation models—to build robust solutions for autonomous robot navigation. Key job responsibilities - Develop and implement robust navigation systems that enable reliable autonomous operation in complex, dynamic indoor environments with static and dynamic obstacles - Build simulation-based and on-device evaluation frameworks with comprehensive benchmarks and metrics for systematic comparison of navigation methods - Conduct sim-to-real transfer experiments, analyzing performance gaps and developing techniques to ensure reliable real-world navigation performance - Collaborate with world model, manipulation, and other teams to ensure seamless integration of navigation capabilities into the full robot system - Stay current with the latest advances in robot navigation, spatial reasoning, and related fields, and apply relevant findings to improve system performance - Mentor fellow scientists and engineers while maintaining strong individual technical contributions About the team Fauna Robotics, an Amazon company, is building capable, safe, and genuinely delightful robots for everyday life. Our goal is simple: make robots people actually want to live and interact with in everyday human spaces. We believe that future won’t arrive until building for robotics becomes far more accessible. Today, too much effort is spent reinventing the fundamentals. We’re changing that by developing tightly integrated hardware and software systems that make it faster, safer, and more intuitive to create real-world robotic products.