Post-Training Research Engineer

About the role

ABOUT BASETEN

Baseten powers mission-critical inference for the world's most dynamic AI companies, like Cursor, Notion, OpenEvidence, Abridge, Clay, Gamma, and Writer. By uniting applied AI research, flexible infrastructure, and seamless developer tooling, we enable companies operating at the frontier of AI to bring cutting-edge models into production. We're growing quickly and recently raised our $1.5B Series F https://www.baseten.co/blog/announcing-our-series-f/, led by Altimeter Capital, Conviction Partners, and Spark Capital. Join us and help build the platform engineers turn to ship AI products.

We are looking for an engineer with strong experience in machine learning and solid foundations in math and computer science to join our growing Post-Training team at Baseten.

THE ROLE

Custom models are instrumental to the success of Baseten customers. By inference volume, the overwhelming majority of traffic at Baseten is to and from models that have been post-trained in some way, whether that be through reinforcement learning, supervised finetuning, a recent technique from the literature, or an in-house research technique from Baseten. The Post-Training team is responsible for the success of our customers’ post-trained models, and we employ a wide array of techniques to produce models that are more efficient and higher quality than even the biggest closed source models for the customer’s specific needs.

Your role as a research engineer is to build the in-house tooling to support all of this. We care about training a wide spectrum of different model architectures with a variety of techniques efficiently and at scale. At times this involves zooming deep into a particular technical topic, but more often it involves working across the stack as a whole - systems-level concepts like Kubernetes, cgroups, storage systems, and networking topologies, as well as PyTorch distributed tensor computation, and GPU kernels.

The Manifesto: https://labs.baseten.co/manifesto

RECENT RESEARCH

- Dense, on-policy or both? https://www.baseten.co/research/dense-on-policy-or-both/

- Repeated kv cache for long-running agents https://www.baseten.co/research/repeated-kv-cache-for-long-running-agents/

- Distillation without the dark – replicating black-box on-policy distillation on Baseten https://www.baseten.co/research/distillation-without-the-dark/

We don’t have a rigid set of skills, but here’s some of what we’re looking for:

- A deep understanding of modern ML techniques and tools for training transformers

- Advanced experience in a tensor/array computation library like PyTorch, TensorFlow, Jax, or similar

- A detailed understanding of transformer training parallelism strategies like data parallelism, sharded data parallelism, tensor parallelism, pipeline parallelism, context parallelism

- The experience and knowledge to profile and improve the performance of a distributed GPU program in PyTorch or a similar library