Description
AWS is hiring a Senior Machine Learning Engineer for its Neuron Distributed Training team to develop, enable, and tune distributed training solutions for AWS Trainium and Inferentia accelerators. The role focuses on building distributed training support for PyTorch and JAX using XLA, the Neuron compiler, and runtime stacks, while optimizing large language models, Stable Diffusion, Vision Transformers, and other model families for peak performance on AWS custom silicon. The position requires a bachelor's degree in computer science or equivalent, at least five years of professional software development experience, and experience with machine learning, distributed training libraries, and model training.
