Sr. Software Engineer- AI/ML, AWS Neuron Distributed Training

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

Инженер в команду Distributed Training Annapurna Labs (AWS) для оптимизации распределённого обучения на кастомных ML-ускорителях Trainium. Работа на стыке ML, HPC и распределённых систем: от PyTorch/JAX и компилятора Neuron до аппаратного уровня.

Задачи

• Оптимизировать производительность распределённого обучения на Trainium: throughput, MFU, время до сходимости
• Работать с PyTorch, JAX, компилятором и рантаймом Neuron для запуска и тюнинга крупномасштабных тренировок
• Запускать новые архитектуры моделей на Trainium: определять недостающие операторы, стратегии шардирования и численные форматы
• Отвечать за стратегии параллелизма: data, tensor, pipeline, expert, context parallelism
• Применять reduced-precision форматы там, где это даёт измеримый выигрыш
• Профилировать end-to-end и определять узкое место: compute, memory, collectives или host overhead
• Формулировать требования к будущей архитектуре Trainium на основе найденных performance gaps
• Контрибьютить в open source фреймворки
• Проектировать и реализовывать бизнес-критичные фичи и менторить инженеров

Требования

• Bachelor's degree in computer science или эквивалент
• 5+ лет профессионального опыта разработки ПО (не стажировки)
• 5+ лет программирования минимум на одном языке
• 5+ лет опыта ведущего проектирования или архитектуры новых и существующих систем (паттерны, надёжность, масштабирование)
• 5+ лет опыта полного SDLC: стандарты кодирования, code review, SCM, сборка, тестирование, эксплуатация
• Опыт ментора, tech lead или руководства инженерной командой
• Опыт в machine learning, data mining, information retrieval, статистике или NLP

Условия

• Полный пакет: sign-on payments и RSU
• Медицинская страховка (medical, dental, vision, prescription), Basic Life & AD&D, Supplemental life, EAP, Mental Health Support, Medical Advice Line, FSA, Adoption and Surrogacy Reimbursement
• 401(k) matching
• Оплачиваемый отпуск и parental leave
• Работа в небольшой agile-команде в культуре обучения и экспериментов
• Коллаборация с applied scientists, system engineers и product managers

Навыки

• PyTorch
• JAX
• TensorFlow
• ML compiler
• Distributed training
• Data parallelism
• Tensor parallelism
• Pipeline parallelism
• Expert parallelism
• Context parallelism
• Reduced precision
• High-performance computing
• Computer architecture
• AWS Trainium
• AWS Neuron
• Python