Principal Engineer: Annapurna Labs Training and Inference , Annapurna Labs, AWS Neuron SDK

Description

At Amazon, Principal Engineers are both technical leaders and hands-on builders. As Amazon's most senior individual contributors, Principal Engineers work on our hardest problems, help establish technical standards, and drive the overall technical architecture, engineering practices, and methodologies for products that delight our customers. On the AWS Neuron SDK team within Annapurna Labs, that mandate is AI model performance on Amazon's custom silicon.

We are looking for a Principal Engineer to drive AI model performance across the lifecycle on AWS Trainium — pre-training, post-training and alignment, and inference serving — with deep hands-on expertise in one of these areas and the drive to influence across all three. Neuron is Amazon's software stack: the layer between customer ML frameworks and Trainium / Inferentia silicon. The responsibilities are broad: you will have the opportunity to drive deeply into model enablement across the path from a model's first pre-training step through aligned post-training to cost- and latency-optimized serving at scale.

You will drive performance for a wide range of AI models, including the leading open-source models and direct collaborations with foundation-model developers, as well as supporting Amazon's own internal model development initiatives. This is a deep hands-on builder role: you will write code, own the forward and backward path, and make binding technical decisions, while influencing direction across the broader Neuron organization, partner teams (EC2, EKS, VPC, Bedrock, HyperPod), and leading ML customers. Annapurna offers full-stack ownership — chip to server to network to compiler to cloud — so you see your work run at customer scale, something that doesn't exist at a standalone chip company or startup. Over time, this person ends up owning the whole process end-to-end.

Key job responsibilities

As the Principal Engineer for Annapurna Labs Training and Inference, you will bring deep hands-on expertise in one of the following areas and influence across all three:

Pre-Training Performance

●Kernel optimization for ML training — forward and backward path — and drive performance out of it, with training collective algorithms (collectives, all-reduce, cross-accelerator coordination) as the primary source of gains.

●Prove top-tier performance even at small scale, and clear the path so customers can land AI models on the chip out of the box. This is a performance role, not an Applied Science role — you partner with scientists but own making it fast, not making it correct.

Post-Training & Alignment

●Supervised fine-tuning, preference optimization, reinforcement learning, and alignment stack that turns pretrained base models into production-grade, aligned models on Trainium — the main lever on delivered model quality and the fastest-growing, least-covered workload class on Trainium today.

●Deliver first-class RLHF/RLAIF and preference methods (PPO, GRPO, DPO/IPO/KTO, reward modeling) with training-plus-inference co-location and convergence safeguards (KL control, reward-scaling, policy-collapse detection), and native performance for the post-training ecosystem (TRL, veRL, NeMo-Aligner, TorchTitan/TorchAO, JAX).

Inference / Serving

●Inference optimization strategy within the Neuron SDK — adapting inference algorithms to Trainium's unique architecture to unlock performance existing industry approaches leave on the table, and maximizing the chip's compute and memory bandwidth to drive cost and latency improvements for large-scale model serving.

●Influence hardware-software co-design by partnering with Annapurna's chip architecture teams so future silicon incorporates lessons from production model serving.

Influence across all three

●Serve as a technical leader on demanding, cross-functional projects spanning pre-training, post-training, and inference; functionally decompose complex problems into simple, straightforward solutions.

●Exert technical influence across multiple teams, mentor senior engineers, and raise the technical bar across the organization.

●Balance technical leadership with strong business judgment to make the right technology choices.

About the team

The AWS Neuron team builds the software stack — compiler, runtime, and ML frameworks — that powers Amazon's Trainium and Inferentia accelerators in EC2 Trn1/2 and Inf1 servers. Our Post-Training group focuses on making Trainium the best place to align and refine frontier models: reinforcement learning (RLHF/RLAIF), preference optimization (PPO, GRPO, DPO), and reward modeling at massive scale. We work end-to-end across a vertically integrated stack — hardware, compiler, runtime, PyTorch/JAX (XLA), and the open-source post-training ecosystem (TRL, veRL, OpenRLHF, NeMo-Aligner, Ray) — turning research techniques into production-grade, high-throughput systems that run efficiently across thousands of nodes.

We are a hands-on, customer-obsessed team that values deep technical ownership and collaboration across chip and software boundaries. You'll work directly with AWS customers and internal science teams to bring up new model families, remove scaling bottlenecks, and ship capabilities that let builders post-train and align large language and multi-modal models faster and more cost-effectively on AWS.

Basic Qualifications

• 10+ years of non-internship professional software development experience

• 7+ years with distributed training or RL systems at scale; strong C/C++ and Python; experience with AI/ML model execution, inference serving, or compiler optimization; and technical leadership/mentorship of engineering teams.

Preferred Qualifications

• Deep hands-on expertise in one of the following, with the ability to influence across all three (not expected to be expert in all): ML model architecture (transformers, MoE, diffusion) and distributed training at scale; post-training methods (SFT, RLHF/RLAIF, DPO/IPO/KTO, reward modeling, PPO/GRPO) with convergence/alignment depth; AI inference optimization and LLM serving on custom accelerators (TPU/GPU/ASIC); ML/post-training framework development (JAX, PyTorch/XLA, TRL, veRL, NeMo-Aligner, TorchTitan/TorchAO); algorithmic optimization (compiler, graph, kernel) and numerical debugging; or adjacent backgrounds (EDA, video pipeline, GPU graphics, HPC).

Amazon is an equal opportunity employer and does not discriminate on the basis of protected veteran status, disability, or other legally protected status.

Los Angeles County applicants: Job duties for this position include: work safely and cooperatively with other employees, supervisors, and staff; adhere to standards of excellence despite stressful conditions; communicate effectively and respectfully with employees, supervisors, and staff to ensure exceptional customer service; and follow all federal, state, and local laws and Company policies. Criminal history may have a direct, adverse, and negative relationship with some of the material job duties of this position. These include the duties and responsibilities listed above, as well as the abilities to adhere to company policies, exercise sound judgment, effectively manage stress and work safely and respectfully with others, exhibit trustworthiness and professionalism, and safeguard business operations and the Company’s reputation. Pursuant to the Los Angeles County Fair Chance Ordinance, we will consider for employment qualified applicants with arrest and conviction records.

Our inclusive culture empowers Amazonians to deliver the best results for our customers. If you have a disability and need a workplace accommodation or adjustment during the application and hiring process, including support for the interview or onboarding process, please visit https://amazon.jobs/content/en/how-we-hire/accommodations for more information. If the country/region you’re applying in isn’t listed, please contact your Recruiting Partner.

The base salary range for this position is listed below. Your Amazon package will include sign-on payments and restricted stock units (RSUs). Final compensation will be determined based on factors including experience, qualifications, and location. Amazon also offers comprehensive benefits including health insurance (medical, dental, vision, prescription, Basic Life & AD&D insurance and option for Supplemental life plans, EAP, Mental Health Support, Medical Advice Line, Flexible Spending Accounts, Adoption and Surrogacy Reimbursement coverage), 401(k) matching, paid time off, and parental leave. Learn more about our benefits at https://amazon.jobs/en/benefits .

USA, CA, Cupertino - 230,100.00 - 311,200.00 USD annually

USA, WA, Seattle - 200,100.00 - 270,600.00 USD annually