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

Amazon is hiring a Principal Engineer: Annapurna Labs Training and Inference , Annapurna Labs, AWS Neuron SDK in Cupertino, United States. It pays $230k-$311k a year and Level rates it ; you can apply on Level.

AI in this role

bedrockpytorchjax

fine-tuning

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

How we rate this

Principal Engineer: Annapurna Labs Training and Inference , Annapurna Labs, AWS Neuron SDK at Amazon rates 99 out of 100 for how much of the daily work is AI. That makes it Builds AI (AI Level 4 of 4). The level is about AI in the job, not seniority.

Classification

Builds AI. The job is building AI systems.

●●●● Builds AI80 to 100
●●●○ Works on AI60 to 79
●●○○ Uses AI40 to 59
●○○○ Little AI0 to 39

Levels come from how often the tools, models and workflows of the role are named in the posting itself. Open the description and count.

Key facts

Principal Engineer: Annapurna Labs Training and Inference , Annapurna Labs, AWS Neuron SDK is an open job at Amazon.
Level rates this role AI Level 4 of 4: Builds AI.
The listed salary is $230k-$311k a year.
The role is based in US, CA, Cupertino.
The job was posted on October 14, 2026.
Seniority: lead.
Applications are handled by Amazon. Apply on Amazon

Prepare for this job

Built from this posting and from similar jobs on Level.

Skills and AI tools this role asks for

Fine-tuningBedrockPyTorchJax

Questions you could be asked

Walk me through fine-tuning a model: what data did you use, and how did you check the result?
Walk me through how you've used Bedrock in your day-to-day work.
What are the limits of PyTorch that you've run into, and how did you work around them?
What's a project where you used Jax hands-on?
How would you decide a model or AI system is ready to ship?

Adapt your resume

List these exact terms on your resume: Fine-tuning, Bedrock, PyTorch, and Jax. An applicant tracking system matches the wording, not the idea.
Attach one line of real, concrete experience to at least one of them. A tool named with nothing behind it rarely survives a human read.
Lead with what you built, trained or shipped. This role is judged on the AI system itself, not the tools around it.

Want an expert to read your CV for this job?

Free. Send your CV and the role you want next. We reply by email within 2 to 4 business days.

Get a free CV review Check your resume with the ATS resume checker

Get new software engineering jobs (Builds AI ●●●●) by email

One email a week with the new software engineering jobs (Builds AI ●●●●), each rated for how much AI is in the work. No recruiter spam, unsubscribe in one click.

Free. One email a week. Unsubscribe in one click.