Senior ML Infrastructure Engineer – Training & Inference

Nimble; San Francisco, Northern (CA, KY); Hybrid; USD 210,000 - 300,000; Full time

Benefits offered by this job

Unlimited Flexible Time Off

Health Insurance

Equity

401k

Referral Bonus

Parking

Job summary

Nimble is hiring a Senior/Staff Software Engineer for its ML Infrastructure team to build and optimize training and inference systems powering autonomous warehouse robots. You will own end-to-end training infrastructure, maximize GPU utilization, and enable researchers to launch new experiments with a single command.

You will collaborate with ML and Robotics teams, scale GPU clusters, and contribute to robust, low-latency pipelines and high-throughput data workflows.

Qualifications

• PhD/masters/BS with equivalent experience in CS or related field.
• 4+ years in infrastructure, distributed systems, ML systems or robotics.
• Experience with Rust, Go, Python, or C++.
• Experience with PyTorch or JAX.
• Strong distributed systems fundamentals, memory management and performance tuning.
• Experience with Kubernetes orchestration and data-heavy distributed jobs.
• Ability to debug across GPU memory, networking, filesystems, and multi-GPU setups.
• Strong cross-functional communication and growth mindset.

Responsibilities

• Design, develop and maintain ML training infrastructure for efficient experiments.
• Build low-latency inference pipelines for production robotics workloads.
• Develop, tune and optimize CUDA kernels.
• Design scalable training platforms with high-throughput data ingestion and distribution.
• Lead design reviews and evaluate technology tradeoffs.
• Review code for quality, tests and performance; mentor engineers.

Skills

Rust

Go

Python

C++

PyTorch

JAX

Kubernetes

Distributed systems

GPU memory management

Education

Bachelor’s, Master’s, or PhD in Computer Science or related field

Tools

PyTorch DDP/FSDP

DeepSpeed

Megatron

NCCL

CUDA

Arrow/Parquet

Job description

Nimble is hiring a Senior/Staff Software Engineer for its ML Infrastructure team to build and optimize training and inference systems powering autonomous warehouse robots. You will own end-to-end training infrastructure, maximize GPU utilization, and enable researchers to launch new experiments with a single command.

You will collaborate with ML and Robotics teams, scale GPU clusters, and contribute to robust, low-latency pipelines and high-throughput data workflows.