Senior ML Engineer, Scalable Training & RL Systems

Apple Inc.; Seattle (WA); On-site; USD 175,000 - 309,000; Full time

Benefits offered by this job

Medical and dental coverage

Retirement benefits

Employee stock purchase plan

Discretionary stock awards

Relocation assistance

Job summary

Apple Inc. in Seattle, WA is seeking a Senior/Staff Machine Learning Engineer to join the Cloud AI Platform team on onsite duties. You will shape technical direction across teams, build training recipes, and deliver scalable infrastructure for reinforcement learning and post-training applications.

You will collaborate with researchers and engineers to translate model requirements into practical training features and optimizations, including long-context MoE, distillation, and multimodal training.

Qualifications

• 8+ years of experience in large-scale model training, including reinforcement learning.
• 5+ years building and operating large-scale model training systems in production.
• Experience with fine-tuning and distillation techniques such as supervised fine-tuning and LoRA.
• Hands-on with Megatron-LM, DeepSpeed, TorchTitan, NeMo, AXLearn, or MaxText.
• Deep expertise in PyTorch or JAX with distributed training features.
• Strong Python or C++ programming and cross-team leadership.
• BS, MS, or PhD in CS/CE or equivalent practical experience.

Responsibilities

• Partner with researchers and ML engineers to develop advanced ML models and training features.
• Develop model distillation and compression approaches to enable smaller models.
• Train models with long-context extension, Mixture-of-Experts, and multimodal inputs.
• Investigate training behavior issues and ensure privacy and security standards.
• Build scalable training frameworks and services on GPUs and TPUs using PyTorch or JAX.
• Create reinforcement learning infrastructure for post-training, including rollout and secure sandboxes.
• Set technical direction across teams and mentor engineers.

Skills

Large-scale model training

Model training systems

Fine-tuning & distillation

PyTorch or JAX

Python or C++

Distributed training

Education

BS/MS/PhD in CS/CE or equivalent

Tools

Megatron-LM

DeepSpeed

TorchTitan

NeMo

AXLearn

MaxText

NCCL

NVLink

InfiniBand

TPU ICI

Job description

Apple Inc. in Seattle, WA is seeking a Senior/Staff Machine Learning Engineer to join the Cloud AI Platform team on onsite duties. You will shape technical direction across teams, build training recipes, and deliver scalable infrastructure for reinforcement learning and post-training applications.

You will collaborate with researchers and engineers to translate model requirements into practical training features and optimizations, including long-context MoE, distillation, and multimodal training.