Senior Software Engineer, AI Infra

Why RoboForce

Responsibilities

• Build and maintain distributed training pipelines leveraging PyTorch, JAX, or equivalent frameworks across multi-GPU and multi-node clusters.
• Build tools for data collection, training, and deployment of neural networks on RoboForce robots.
• Architect robust cloud-native and on-prem GPU infrastructure across multi-cloud environments.
• Build high-throughput data workflows to support large-scale dataset ingestion, versioning, and distributed storage.
• Optimize end-to-end training performance: CPU–GPU transfers, NVMe caching, I/O pipelines, containerized runtime environments, and CUDA-level optimizations.
• Integrate training artifacts into on-robot inference stacks.

Requirements

• Bachelor’s or Master’s degree in Computer Science or related field with 5+ years of experience.
• Strong proficiency with C++, Python, and ML frameworks (e.g., PyTorch, JAX).
• Deep experience with at least one major cloud provider (GCP, AWS, Azure) and familiarity with Kubernetes, containerization, and GPU machine provisioning.
• Strong understanding of SQL and NoSQL data stores (Postgres, MySQL, BigQuery, ElasticSearch, Redis).
• Requires 5 days/week in-office collaboration with the teams.

Bonus Qualifications

• Expertise in profiling and optimizing CPU-GPU interactions.
• Experience scaling neural network training jobs and GPU programming with CUDA.
• Proven ability to develop annotation and dataset management tools.

Benefits

• Competitive stock options/equity programs.
• Health, dental, and vision insurance, 401(k) plan.
• Visa sponsorship and green card support for qualified candidates.
• Lunches and dinners, a fully stocked kitchen, and regular team-building events.