Software Engineer (Training Infrastructure)

• As a Software Engineer at on the Training Infrastructure team, you’ll architect and lead development of our training platform, supporting top tier research engineers and model developers
• You’ll make key technical decisions for the infrastructure enabling developers to deploy, scale, and monitor their workloads with high performance and reliability
• You’ll own scheduling, storage, networking, reliability, and observability of technical systems in the training stack
• Design and architect scalable infrastructure systems for our ML training platform (e.g. scheduling, storage, and networking)
• Partner closely with developers and research engineers to translate complex training requirements into technical solutions
• Design and architect a global training scheduler
• Design and architect reinforcement learning systems and continuous learning pipelines
• Drive long-term improvements to improve reliability of systems and velocity of development
• Partner closely with SRE and Capacity teams to unlock state of the art training infrastructure
• Make critical architectural decisions balancing performance with system reliability
• Lead technical discussions and mentor junior engineers on infrastructure best practices
• Contribute to long-term technical strategy and infrastructure roadmap

Benefits

• Remote-first work environment. The Baseten team is welcome to work from wherever they want; fully remote, in our San Francisco office, or a mix of both. Today, our team (including our founding team) is spread across the United States, Canada, and Armenia. We provide a $1,000 stipend for you to make your home-office comfortable and productive

• Regular in-person team summits. We get together as a team three times a year to plan, workshop, and most importantly, get to know each other better

• Unlimited PTO. We ask that everyone take at least 4 weeks of vacation. And we have a company-wide break between Christmas and New Year’s Day

• Full healthcare coverage. Medical, dental and vision insurance for you and your family

• Paid parental leave. 16-weeks fully paid parental leave (adoptive and non-birth parents included) and flexibility with schedules while returning to work

• Company-sponsored 401(k) for you to contribute to

• Learning and development budget. We encourage you to take classes, attend conferences, and invest in your craft and we’ll cover expenses to make it happen

• 5+ years of experience

• Proficiency in Go, with

• Experience with ML/AI workloads and MLOps platforms

• Advanced understanding of distributed systems concepts and performance tuning

• Proven experience designing observability systems

• Bachelor’s degree or higher in Computer Science or related field

• Deep expertise with Kubernetes in production environments

• Experience with distributed storage systems

• Python experience a plus

• Extensive experience with major cloud providers (AWS, GCP) and neo-cloud providers (Crusoe, DigitalOcean, Nebius)

• Experience with workload orchestration platforms like Temporal or Airflow

• Familiarity or experience with the open source training stack and frameworks (NCCL, PyTorch, Megatron, NemoRL, VeRL, Axolotl, HF Trainer) and distributed training techniques (FSDP, DeepSpeed)

• Experience developing AI products, tooling, or agents