A

Apple

Designs phones, computers and their operating systems

Artificial Intelligence and Machine Learning Senior Machine Learning/Reinforcement Learning Trainin…

Applied Machine LearningConfirmed live in the last 24 hours
Full-timeZürich, SwitzerlandNo salary listedIn-person

About the Job

Requirements

  • PhD or Master of Science in Computer Science, Computer Engineering, or a closely related field.
  • Hands-on experience designing, building, or maintaining large-scale machine learning training infrastructure.
  • Strong proficiency with PyTorch or JAX and experience running training workloads on GPUs or TPUs.
  • Solid understanding of distributed systems concepts, including parallelism strategies, fault tolerance, and synchronization.

Responsibilities

  • Design, build, and scale systems that enable large-scale reinforcement learning for foundation models.
  • Focus on TPU-based training with JAX and develop robust, high-performance reinforcement learning pipelines that support distributed actor/learner architectures, efficient experience replay, and large-scale environment execution.
  • Work across the full stack of reinforcement learning 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 so research teams can iterate quickly and explore more complex reinforcement learning environments and models.
  • Improve the scalability, throughput, and stability of reinforcement learning experiments.

Desired Qualifications

  • Practical experience developing or optimizing training loops, reinforcement learning pipelines, or large-scale model-training frameworks.
  • Strong software engineering skills in Python, emphasizing reliability, debuggability, and high-performance execution.
  • Deep experience with PyTorch or JAX internals, XLA, debugging, and performance profiling on GPU or TPU architectures.
  • Expertise in distributed reinforcement learning 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 machine learning jobs, including I/O, input pipelines, kernel performance, memory, and compilation.
  • Familiarity with reinforcement-learning-specific infrastructure requirements, such as actor/learner architectures, experience replay systems, and large-scale environment execution.
  • Strong software engineering practices, including code quality, design reviews, testing, observability, and continuous integration/continuous delivery.
  • Experience working with cloud-scale clusters or specialized accelerators, including TPU v5/v6, GPUs, or custom hardware.
  • Contributions to machine learning frameworks, distributed training libraries, or high-performance computing systems.

About the Company

Apple

Designs phones, computers and their operating systems

Apple designs and sells consumer electronics, from iPhone smartphones and Mac computers to tablets and wearable devices. Its operating systems manage apps and built-in functions for communication and work, while iCloud stores files and synchronizes data across devices. Customers buy additional digital services through subscriptions such as Apple Music and Apple TV, and developers distribute software through the App Store. Apple earns money from hardware sales, subscriptions and digital transactions, selling directly to customers and through retailers and mobile carriers around the world. It was founded in 1976 and has its headquarters in Cupertino, California.

Company size
10,001+
Company stage
IPO
Headquarters
Cupertino, CA
Founded
1976

Recently Posted Jobs

Explore more opportunities at Apple

See more jobs