Senior / Staff ML Training Optimization Engineer

Highlights

• Build standardized distributed training frameworks for research and production.
• Drive training toward new levels of stability and efficiency at scale.
• Comprehensively profile model runtime and memory to pinpoint bottlenecks.

Ideal candidate

Collaborative, open-minded engineer with 4+ years experience who loves solving hard problems and optimizing large-scale training. Passionate about self-driving technology and helping researchers use resources efficiently.

Requirements

• MS/PhD or Bachelor's degree with minimum 4 years industry experience in Computer Science, Robotics or similar field
• Solid coding proficiency in Python, C++ or Rust
• Experience with deep learning frameworks such as PyTorch or JAX
• Skilled in profiling CPU and GPU code using PyTorch Profiler and NVIDIA Nsight
• Open-minded and collaborative team player with willingness to help others
• Passionate about self-driving technologies, solving hard problems and creating innovative solutions
• Nice to have: experience identifying when custom CUDA kernels are needed and implementing them
• Nice to have: experience with Bazel in a monorepo environment and integrating third-party packages
• Nice to have: experience with Kubernetes-based training platforms