Staff Machine Learning Engineer (Infra)(TLM)

Location: Onsite in Mountain View, California, United States

Salary: $251,000–$310,000 USD per year, plus eligibility for an annual bonus, equity incentive plan, and company benefits.

Company

Waymo develops autonomous driving technology and operates a fully autonomous ride-hailing service powered by the Waymo Driver.

What you will do

• Provide technical leadership for large-scale machine learning model architectures used in autonomous vehicles.
• Build scalable systems for training and fine-tuning models that evaluate driving behaviors.
• Design distributed systems supporting dataset generation, model training, and evaluation across the ML lifecycle.
• Operate and optimize models assessing data from Waymo’s autonomous vehicle fleet.
• Guide architectural decisions and technical direction for complex systems.
• Collaborate across functions to translate product and business goals into measurable technical deliverables.

Requirements

• M.S. or Ph.D. in Computer Science, Machine Learning, Artificial Intelligence, or a related technical field, or equivalent practical experience.
• 7+ years of professional software engineering experience, including at least 3 years building, scaling, training, deploying, and optimizing machine learning infrastructure.
• Experience with machine learning tooling and frameworks such as PyTorch, JAX, TensorFlow, Ray, or similar, including their user-facing APIs and internal workings.
• Strong expertise in distributed training, gradient sharding, optimization strategies, and ML accelerator profiling.
• Deep understanding of state-of-the-art models, including autoregressive transformers.
• Strong leadership skills and experience providing technical direction across cross-functional organizations.

Nice to have

• 10+ years of software engineering experience, including at least 5 years in machine learning infrastructure.
• Experience in autonomous vehicles, robotics, or complex simulation environments.
• Knowledge of reinforcement learning techniques for fine-tuning large models using human feedback or preferences.
• Experience with large-scale simulation platforms, AI evaluation metrics, and cross-organizational technical leadership.
• Excellent communication skills for explaining complex technical concepts.

Culture & Benefits

• Work on machine learning systems supporting autonomous driving technology and large-scale vehicle evaluation.
• Eligibility for a discretionary annual bonus program and equity incentive plan.
• Generous company benefits program, subject to eligibility requirements.