← All roles

AIML - Senior ML/RL Training Infrastructure Engineer, AFM

Apple· Machine Learning and AI
Apply on company site
Hybrid (Zürich)full-timePosted 6mo ago

About the role

Senior ML/RL training infrastructure engineer for Apple’s Core Foundation Models team, focusing on large-scale RL training infrastructure, TPU-based training with JAX, and distributed systems to enable research and production-scale experiments across multiple Apple engineering hubs.

Responsibilities

  • Design, build, and scale systems that enable large-scale reinforcement learning for Apple’s foundation models.
  • Focus on TPU-based training with JAX, developing robust, high-performance RL pipelines that support distributed actor/learner architectures, efficient experience replay, and large-scale environment execution.
  • Work across the full stack of RL training systems—from low-level performance tuning and compiler optimization to cluster-level orchestration and resource management.
  • Ensure training pipelines are efficient, reliable, reproducible, and observable, enabling research teams to iterate quickly and explore more complex RL environments and models.
  • Collaborate with teams across Apple’s engineering hubs—including New York, Seattle, and Cupertino—to advance tooling and systems for large-scale model training.

Qualifications

  • PhD or MSc in Computer Science, Computer Engineering or a closely related field.
  • Hands-on experience designing, building, or maintaining large-scale ML training infrastructure.
  • Strong proficiency with PyTorch or JAX and experience running training workloads on GPUs/TPUs.
  • Solid understanding of distributed systems concepts (parallelism strategies, fault tolerance, synchronization).
  • Practical experience developing or optimizing training loops, RL pipelines, or large-scale model-training frameworks.
  • Strong software engineering skills in Python, with emphasis on reliability, debuggability, and high-performance execution.
  • Deep experience with PyTorch/JAX internals, XLA, debugging and performance profiling on GPU/TPU architectures.
  • Expertise in distributed RL training patterns, including actor/learner architectures, experience replay, and parallel environment execution.
  • Experience building training services, orchestration tools, or automated pipelines for large-scale experiments.
  • Proven success diagnosing bottlenecks in large-scale ML jobs (I/O, input pipelines, kernel performance, memory, compilation).
  • Familiarity with RL-specific infrastructure requirements (e.g., actor/learner architectures, experience replay systems, large-scale environment execution).
  • Strong software engineering practices: code quality, design reviews, testing, observability, CI/CD.
  • Experience working with cloud-scale clusters or specialized accelerators (TPU v5/v6, GPU, custom hardware).
  • Contributions to ML frameworks, distributed training libraries, or high-performance computing systems.
  • Excellent communication and collaboration skills for working with research and engineering partners.

Skills

Required:PhD or MSc in Computer Science/Computer Engineering, or closely related fieldPyTorch or JAXDistributed systemsPythonRL training patterns (actor/learner, experience replay, environment execution)Training pipelines and orchestrationPerformance optimization & profilingCommunication & collaboration
Nice to have:XLACloud-scale clusters / accelerators

Benefits

  • Commitment to fair and equal opportunity employment.
  • Accessibility and accommodations in the workplace.
Last updated Jun 26, 2026. Interested? Apply on the company site.

About Apple

Discover the innovative world of Apple and shop everything iPhone, iPad, Apple Watch, Mac, and Apple TV, plus explore accessories, entertainment, and expert device support.

Discover the innovative world of Apple and shop everything iPhone, iPad, Apple Watch, Mac, and Apple TV, plus explore accessories, entertainment, and expert device support.

AirPods Max 2Apple ArcadeApple CardApple Fitness+Apple MusicApple TVApple TV+Apple Watch Series 11
Company profileWebsite

More open roles at Apple