Physical Intelligence
ML Infra Engineer, Modeling
San Francisco · On-site · Full-time · Posted Sep 1, 2026
About the job
Build and optimize infrastructure and core training code for large-scale JAX-based model training and inference across GPU and TPU clusters.
What you'll do
- Design and maintain training and inference systems, including scheduling, checkpointing, and metrics/logging.
- Scale distributed training and optimize memory use, device utilization, throughput, and synchronization.
- Build tools for experiment iteration and collaborate with researchers on training infrastructure and core model code.
What they ask for
- Strong software engineering fundamentals and experience building ML training infrastructure or internal platforms.
- Hands-on large-scale training experience in JAX or PyTorch.
- Familiarity with distributed training, multi-host setups, data loaders, and evaluation pipelines.
- Experience managing training workloads on cloud platforms and ability to debug and optimize performance bottlenecks.
Nice to have
- Deep ML systems background, including training compilers, runtime optimization, or custom kernels.
- Experience with GPU/TPU performance tuning, robotics, multimodal models, or large-scale foundation models.
Tools
JAX, PyTorch, SLURM, Kubernetes, GCP TPU/GKE, AWS
AI focus
Supports research and development of foundation models and learning algorithms for robots and physically actuated devices; contributes to model training code for new architectures, modalities, and evaluation metrics.
Model-written summary · the original posting is authoritative