Search
Go
ML Infra Engineer, Modeling
Summary by AIEngineer.careers
ML Infra Engineer, Modeling is a AI Infrastructure Engineer role (full-time). with Physical Intelligence. in SAN FRANCISCO, US. Imported listing (source: zerogtalent.com). Apply on the employer's site (zerogtalent.com).
Imported listing (source: zerogtalent.com) · Apply on zerogtalent.com
SAN FRANCISCO · US
—
Full-time
SENIOR
The sections below reproduce the third-party job description for reference. AIEngineer.careers does not write or control this text.
Imported job description
Sourced from zerogtalent.com
Physical Intelligence is seeking an ML Infrastructure Engineer to design and optimize large-scale model training systems. You will own critical infrastructure for distributed training across TPU and GPU clusters, build efficient JAX pipelines, and collaborate closely with researchers to translate experiments into production-grade runs. This hands-on role sits at the intersection of machine learning, software engineering, and scalable infrastructure.
Responsibilities
- Design and maintain large-scale model training and inference infrastructure
- Scale distributed JAX-based training across TPU and GPU clusters
- Profile and optimize memory usage, device utilization, and throughput
- Build abstractions for experiment launching, monitoring, and debugging
- Partner with researchers to translate needs into scalable infra capabilities
Requirements
- Strong software engineering fundamentals with experience building ML training infrastructure
- Hands-on large-scale training experience in JAX (preferred) or PyTorch
- Familiarity with distributed training, multi-host setups, and evaluation pipelines
- Experience managing workloads on cloud platforms like GCP, AWS, SLURM, or Kubernetes
Physical Intelligence
San Francisco · US · 302+ employees
Physical Intelligence is a research and development company focused on building general-purpose foundation models and learning algorithms to power robots and other physical devices. Its goal is to bring AI into the physical world by enabling robots to perform a wide variety of tasks.
