Software Development Engineer II, ML Acceleration

AWS Neuron is the complete software stack for the AWS Inferentia (Inf1/Inf2) and Trainium (Trn1), our cloud-scale Machine Learning accelerators. This role is for a machine learning engineer in the Inference team for AWS Neuron, responsible for development, enablement and performance tuning of a wide variety of ML model families, including massive-scale Large Language Models (LLM) such as GPT and Llama, as well as Stable Diffusion, Vision Transformers (ViT) and many more.

The ML Inference team works side by side with chip architects, compiler engineers and runtime engineers to create, build and optimize distributed inference solutions with Trainium/Inferentia instances. Experience with training and optimizing inference on these large models using Python/C++ is a must. Model parallelization, quantization, memory optimization - vLLM, DeepSpeed and other distributed inference libraries are central to this role, and extending all of them for the Neuron based system is key.

Key responsibilities include:

• Leading efforts to build and achieve the best distributed training and inference performance of PyTorch, JAX, TensorFlow with XLA and other advanced frameworks on Neuron stacks.
• Optimizing models to ensure the highest performance and maximize efficiency on custom AWS Trainium and Inferentia silicon and the Trn1, Inf1/2 servers.
• Strong software development (Python and C++) and Machine Learning knowledge are critical to this role.

This position offers an opportunity to work on cutting-edge ML accelerator technology and contribute to the development of AWS's cloud-scale machine learning infrastructure.