Sr. Software Engineer- AI/ML, AWS Neuron Distributed Training
Summary
Builds and optimizes AWS Neuron’s distributed training stack for PyTorch/JAX models on Trainium, enabling frontier-scale GenAI workloads.
Description The Annapurna Labs team at Amazon Web Services (AWS) builds AWS Neuron, the software development kit used to accelerate deep learning and GenAI workloads on AWS Trainium, Amazon's custom machine learning accelerator. Neuron includes an ML compiler, runtime, collectives library, and application framework that integrate with PyTorch and JAX, so customers can train frontier-scale models on Trainium without rewriting their stack. The Distributed Training team is at the forefront of trai…