Senior Machine Learning Engineer, AI Infrastructure, Autonomy

Job Description

Rivian is seeking a full-time Machine Learning Engineer to establish a cutting-edge AI infrastructure for training large foundation models and accelerating model training/inference for autonomous driving.

Requirements

• PhD in CS/CE/EE or equivalent industry experience
• Deep knowledge of PyTorch
• Experience with Cuda or Triton language for writing custom ops
• Knowledge of model training framework (e.g., PyTorch Lightning)
• In-depth knowledge of transformer architecture and ways to accelerate training and inference of transformer models
• Experience in large scale distributed training of models
• Profiling model and improving model training and inference speed

Responsibilities

• Design, train, and deploy large deep learning models leveraging labeled and unlabeled data
• Improve ecosystem for training infrastructure and deployment pipeline
• Work with a team to push boundaries of machine learning models for autonomous driving

Benefits

• Robust medical/Rx, dental and vision insurance packages for full-time and part-time employees, their spouse or domestic partner, and children up to age 26.