Software Engineer, ML Training Performance

Автоматический пересказ вакансии на русском. Полный текст на языке оригинала и форма отклика — на сайте работодателя.

Инженер в ML Platform команду Zoox для оптимизации производительности обучения крупномасштабных моделей (VLM, VLA, foundation models) для автономного вождения. Работа с распределённым обучением, RL и GPU-кластерами.

Задачи

• Проектировать, внедрять и эксплуатировать техники оптимизации производительности ML-обучения: распределённое обучение (data/tensor/pipeline parallelism), mixed-precision, gradient checkpointing, GPU-ядра
• Работать над RL-фреймворком и оптимизацией производительности RL для foundation-моделей вождения
• Профилировать и отлаживать узкие места пропускной способности обучения на уровне compute, компилятора, фреймворка и кластерной сети
• Тесно сотрудничать с ML-исследователями, data-инженерами и hardware-инженерами для определения требований и архитектурных решений
• Реализовывать кастомные ядра и бенчмарки для отслеживания и валидации оптимизаций обучения

Требования

• 4+ года общего опыта, включая 2+ года работы с платформами крупномасштабного обучения моделей
• Опыт с фреймворками распределённого обучения (Megatron-LM, TorchTitan) и эффективным использованием GPU/accelerator-кластеров
• Опыт с RL-фреймворками (Verl/MILES) и RL-алгоритмами (GRPO, PPO)
• Опыт с библиотеками коллективных коммуникаций (например, NCCL) и оптимизацией перекрытия communication/compute в multi-node обучении
• Опыт использования инструментов профилирования (NVIDIA Nsight, PyTorch Profiler, XLA/JAX profiling)
• Владение Python или C++
• Опыт с ML-компиляторами, performance-инструментами или hardware/software co-design (GPU/TPU/accelerator архитектуры)

Условия

• Зарплата, Amazon RSUs и Zoox Stock Appreciation Rights
• Возможен sign-on bonus
• Оплачиваемый отпуск (больничный, отпуск, bereavement), неоплачиваемый отпуск
• Медицинская страховка, страхование долгосрочного ухода, краткосрочная и долгосрочная страховка по нетрудоспособности, страхование жизни

Навыки

• Python
• C++
• Megatron-LM
• TorchTitan
• Verl
• MILES
• GRPO
• PPO
• NCCL
• NVIDIA Nsight
• PyTorch Profiler
• XLA
• JAX
• GPU
• TPU
• distributed training
• mixed-precision
• gradient checkpointing
• RL