Sr. Software Engineer- AI/ML, AWS Neuron Distributed Training

Annapurna Labs, now fully integrated with AWS after its 2015 acquisition, is seeking a Senior Machine Learning Engineer for their Distributed Training team. This role focuses on AWS Neuron, the complete software stack for AWS Trainium and Inferentia cloud-scale Machine Learning accelerators. The position involves working with cutting-edge ML technologies, including Large Language Models like GPT and Llama, as well as Stable Diffusion and Vision Transformers. The role requires expertise in distributed training libraries such as FSDP, Deepspeed, and Nemo, while working closely with chip architects and compiler engineers. The team values knowledge-sharing, mentorship, and career growth, offering opportunities to work on complex technical challenges that push the boundaries of cloud computing. AWS provides a comprehensive compensation package, including competitive base pay, equity, and extensive benefits, while fostering an inclusive culture that celebrates diversity and work-life harmony. This position represents an opportunity to shape the future of machine learning infrastructure at one of the world's leading cloud platforms.