Software Dev Engineer II, Stores Foundational AI -SFAI

Key job responsibilities

• Design and implementation of a stable and efficient training system for model training and reinforcement learning that scale to various of model sizes and architecture.
• Collaborate with other talented applied scientists and engineers to improve training efficiency and reliability that accelerates innovation.
• Design and implement scalable data infrastructure: that handle Amazon-scale data ingestion, processing, and delivery across different training and evaluation stages;
• Quickly learn and adopt state-of-the-art technologies and algorithms in the field of Generative AI.

A day in the life

• Design and build end-to-end RL post-training pipelines (rollout → reward → optimization) at cluster scale
• Improve RL training stability (PPO / GRPO / RLOO) by monitoring and tuning key metrics such as reward, KL divergence, and policy stability
• Optimize RL post-training efficiency (GPU utilization, batching, sequence packing, async rollouts)
• Partner with research scientists to translate new RL algorithms into scalable, production-ready systems
• Profile and eliminate bottlenecks across compute, networking, and storage
• Build observability systems for training dynamics, system health, and experiment tracking
• Collaborate cross-functionally to run experiments, iterate quickly, and unblock research progress
• Contribute to system design and long-term technical roadmap

About the team

The SFAI Training Infrastructure team builds a unified platform for large-scale LLM training, supporting the full lifecycle from pretraining to fine-tuning and RL post-training. We focus on solving hard system challenges at the intersection of distributed systems and machine learning, building a platform that is:

• Scalable — Efficiently train modern model architectures across large-scale compute environments
• Reliable — Enable long-running jobs through fault tolerance, monitoring, and automated recovery
• Efficient — Maximize hardware utilization and throughput through system-level optimizations
• Simple and Unified — Provide a consistent, config-driven interface across models and workflows

Basic Qualifications

• 3+ years of non-internship professional software development experience
• 2+ years of non-internship design or architecture (design patterns, reliability and scaling) of new and existing systems experience
• Experience programming with at least one software programming language
• Knowledge of Machine Learning and LLM fundamentals, including transformer architecture, training/inference lifecycles, and optimization techniques

Preferred Qualifications

• Knowledge of ML frameworks including JAX, PyTorch, vLLM, SGLang, Dynamo, TorchXLA, and TensorRT
• Knowledge of system performance, memory management, and parallel computing principles
• Experience with CUDA/C++/Kernel development