Software Engineer- AI/ML, AWS Neuron

AWS Neuron is the complete software stack for the AWS Inferentia and Trainium cloud-scale machine learning accelerators and the Trn1 and Inf1 servers that use them. This role is for a senior software engineer in the Machine Learning Applications (ML Apps) team for AWS Neuron. The role is responsible for development, enablement, and performance tuning of various ML model families, including large language models like GPT2 and GPT3, stable diffusion, and Vision Transformers.

Key responsibilities include:

• Leading efforts to build distributed training and inference support into PyTorch, TensorFlow, and JAX using XLA and the Neuron compiler and runtime stacks.
• Tuning models for highest performance and efficiency on AWS Trainium and Inferentia silicon and the TRn1 and Inf1 servers.
• Working closely with chip architects, compiler engineers, and runtime engineers to create, build, and tune distributed training solutions with Trn1.

The role requires strong software development skills and ML knowledge. Experience with Python, FSDP, Deepspeed, and other distributed training libraries is essential.

AWS values work-life balance and offers flexibility in working hours. The team embraces diversity and inclusion, with employee-led affinity groups and ongoing learning experiences. Career growth and mentorship opportunities are prioritized, with projects assigned to help team members develop into well-rounded professionals.

This position offers a competitive compensation package, including base pay ranging from $129,300 to $223,600 per year depending on the geographic market, as well as potential equity, sign-on payments, and a full range of benefits.