Principal ML Systems Engineer - Physical AI Training

Neolix; Sunnyvale (CA); On-site; USD 180,000 - 260,000; Full time

Benefits offered by this job

Premium health insurance

Dental insurance

Vision insurance

401(k)

Job summary

Neolix US, the Silicon Valley Physical AI research team of Neolix, seeks an ML Systems owner to build a scalable training platform for video and multimodal driving data.

You will lead multi-node training execution, profiling of compute and memory, and robust experiment workflows, partnering with security and legal teams to ensure compliant, repeatable research that lowers time and cost per validated experiment.

Qualifications

• Experience at a leading autonomous-driving company.
• Production distributed ML training infrastructure ownership and high-scale runtimes.
• Strong Python and C++, with deep knowledge of PyTorch/JAX execution, data loaders, memory pressure and communications.
• Track record of performance profiling, root-cause analysis and cost improvements.

Responsibilities

• Lead training platform architecture and multi-node training execution and observability.
• Profile compute, memory, interconnects and accelerators; implement fixes for performance.

Skills

Python

C++

PyTorch/JAX

Distributed training

Profiling & debugging

Storage & data movement

Tools

PyTorch

JAX

Distributed training frameworks

CUDA

Job description

Neolix US, the Silicon Valley Physical AI research team of Neolix, seeks an ML Systems owner to build a scalable training platform for video and multimodal driving data.

You will lead multi-node training execution, profiling of compute and memory, and robust experiment workflows, partnering with security and legal teams to ensure compliant, repeatable research that lowers time and cost per validated experiment.