Software Engineer- AI/ML, Amazon Neuron Training

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

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

Задачи

• Реализовывать и настраивать компоненты стека распределённого обучения для крупномасштабного обучения, post-training и RL на инстансах Trainium
• Работать с PyTorch и стеком Neuron
• Участвовать в реализации стратегий параллелизма: data, tensor и pipeline parallelism
• Применять reduced-precision форматы под руководством старших коллег
• Профилировать нагрузки и определять узкие места: compute, memory, collectives или host overhead
• Совместно с инженерами compiler, runtime и collectives доводить исправления
• Создавать и поддерживать внутренний инструментарий, бенчмарки и тесты для воспроизводимости performance-работы

Требования

• Bachelor's degree или выше в computer science или эквивалент
• 2+ года опыта проектирования или архитектуры новых и существующих систем (design patterns, reliability, scaling)
• 3+ года опыта полного цикла разработки ПО в продакшене
• 3+ года опыта с одним из языков программирования: Python, C/C++ или аналогичным
• Понимание основ LLM/transformer: attention mechanisms, autoregressive decoding, KV-cache, формы параллелизма

Условия

• Гибкий график работы и поддержка work-life balance
• Комплексные бенефиты: медицинская страховка (medical, dental, vision), 401(k) matching, оплачиваемый отпуск, parental leave
• Sign-on payments и restricted stock units (RSUs)
• Инклюзивная культура и employee-led affinity groups

Навыки

• Python
• C
• C++
• PyTorch
• JAX
• AWS Neuron
• Trainium
• LLM
• Transformer
• Distributed Training
• Data parallelism
• Tensor parallelism
• Pipeline parallelism
• RLHF
• PPO
• GRPO
• Performance engineering
• Profiling