ML Infrastructure Engineer

6 months ago

Foster City, CA, USAMid Level

H1B Sponsor

Responsibilities

Maintain an effective research codebase optimized for fast iteration and correctness.
Own model-training infrastructure for job scheduling, checkpointing, metrics, and logging.
Scale distributed training across GPU clusters and enable larger models through sharding, activation checkpointing, and memory optimization.
Profile and optimize GPU utilization, memory usage, and training throughput.
Build low-latency inference pipelines for real-time robot control and optimize inference performance through quantization, distillation, and model compilation.
Collaborate with researchers and roboticists to translate research needs into reliable software and infrastructure.
Design high-throughput pipelines for ingesting, validating, and transforming multimodal robot data.
Build storage systems and metadata indexing for large-scale dataset management.
Optimize dataloaders, sharding, and prefetching to reduce the time from data arrival to model training.
Build research tools for debugging, visualization, and experiment analysis.

Strong software engineering and systems fundamentals.
Experience building distributed systems or large-scale data pipelines.
Hands-on experience with ML training infrastructure, ideally PyTorch.
Ability to reason about performance, memory, I/O, and GPU utilization.
Experience managing training workloads with SLURM, Kubernetes, or similar tools.
Ability to design, build, operate, and iterate on systems end-to-end.
Willingness to work closely with researchers and unblock fast-moving projects.
Nice to have experience with robotics data pipelines or multimodal models.
Nice to have a background in VLAs, Video Generation architectures, or robot-learning systems.
Nice to have deep ML systems experience with training compilers, custom kernels, or runtime optimization.
Nice to have hands-on GPU performance-tuning experience.
Nice to have experience with Protobuf, FlatBuffers, or MCAP.

51-200 employees

The helpful home robot company