About the job
The Annapurna Labs team at Amazon 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.
Responsibilities
You will lead efforts to optimize distributed training performance on Trainium, with a primary focus on training throughput, model FLOPs utilization, and time to convergence across the Neuron software stack."You will work across PyTorch, JAX, and the Neuron compiler and runtime to enable and tune large-scale training workloads on the latest Trainium instances.\
Qualifications
Minimum
Bachelor's degree in computer science or equivalent\n5+ years of non-internship professional software development experience\n5+ years of programming with at least one software programming language experience\n5+ years of leading design or architecture (design patterns, reliability and scaling) of new and existing systems experience\n5+ years of full software development life cycle, including coding standards, code reviews, source control management, build processes, testing, and operations experience\nExperience as a mentor, tech lead or leading an engineering team\nExperience in machine learning, data mining, information retrieval, statistics or natural language processing
Preferred
Master's degree in computer science or equivalent\nExperience in computer architecture\nPrevious software engineering expertise with Pytorch/Jax/Tensorflow, Distributed libraries and Frameworks, End-to-end Model Training.