AI/ML Software Engineer - Distributed Training HPC
Listed on 2026-10-07
-
Software Development
Software Engineer, Machine Learning/ ML Engineer, AI Engineer (Applied/Software)
Amazon’s Annapurna Labs team within AWS Neuron Training seeks a Software Engineer to optimize distributed training on Trainium, tackling throughput across the Neuron stack and collaborating with PyTorch and JAX integrations. You will own parallelism strategies and help land improvements across compiler, runtime, and collectives teams.
The role emphasizes performance tuning, end-to-end profiling, and upstream contributions to open source frameworks used by customers training frontier models on
We are seeking a motivated AI/ML Software Engineer - Distributed Training HPC to join Amazon in Cupertino, CA, United States.
The position is based in Cupertino, CA, United States.
This opportunity is part of our work in IT & Technology.
The advertised compensation is 165..
We aim to respond to suitable candidates as soon as possible.
Full responsibilities and requirements are described in the listing above.
(If this job is in fact in your jurisdiction, then you may be using a Proxy or VPN to access this site, and to progress further, you should change your connectivity to another mobile device or PC).