Post-Training Research Engineer
Hybrid (San Francisco)$200k - $275kFull-timePosted 7mo ago
About the role
Engineering role focusing on post-training models, building in-house tooling to support post-trained models, and optimizing training and inference efficiency at scale across the full stack.
Responsibilities
- Build in-house tooling to support post-training models and related workflows
- Work across the stack including systems, Kubernetes, cgroups, storage, networking, PyTorch distributed computation, and GPU kernels
- Profile and improve performance of distributed GPU programs in PyTorch or similar libraries
- Perform roofline analysis on transformer training setups
- Collaborate with researchers to derive specifications from messy problems and execute on them
- Familiarity with HPC and distributed computing platforms (Slurm, Ray, Kubernetes, Dask)
- Understand cluster networking technologies (Infiniband, RoCE, GPUDirect)
- Solid fundamentals in OS concepts, containerization, and networking protocols
- Ask difficult questions about approach, assumptions, and tooling choices
Qualifications
- A deep understanding of modern ML techniques and tools for training transformers
- Advanced experience with tensor/array computation libraries (PyTorch, TensorFlow, Jax, or similar)
- Knowledge of transformer training parallelism strategies (data, sharded data, tensor, pipeline, context parallelism)
- Ability to profile and improve distributed GPU performance in PyTorch or similar
- Capability to perform roofline analysis on transformer training setups
- Willingness to tackle messy problems and collaborate with researchers to derive specs
- Familiarity with HPC and distributed platforms (Slurm, Ray, Kubernetes, Dask)
- Familiarity with cluster networking (Infiniband, RoCE, GPUDirect)
- Solid fundamentals in operating systems, containerization, and networking protocols
- Creative mindset and willingness to question tooling choices and assumptions
Skills
Required:Transformer training and optimizationDistributed training frameworksParallelism strategiesProfiling and performance optimizationRoofline analysisOperating systems and containerization
Nice to have:HPC and distributed systemsCluster networking
Benefits
- Competitive compensation, including meaningful equity
- (U.S. only) 100% coverage of medical, dental, and vision insurance for employee and dependents
- Flexible PTO policy including company wide Winter Break (office closed Christmas Eve to New Year’s Day)
- Paid parental leave
- Fertility and family-building stipend through Carrot
- (U.S. only) Company-facilitated 401(k)
- Exposure to a variety of ML startups for learning and networking opportunities
Last updated Oct 21, 2026. Interested? Apply on the company site.
About Baseten
Inference is everything
Effortlessly serve optimized open source & custom models on the fastest, most reliable model delivery network.
AICloud InfrastructureDeveloper ToolsInference PlatformMachine LearningModel DeploymentModel Runtimesai-infrastructure-and-development
More open roles at Baseten
- Partner Enablement Lead5d agoHybrid (San Francisco), Hybrid (New York) · G&A
- Engineering Manager Runtime - Storage5d agoHybrid (San Francisco) · EPD
- Sales Manager - Enterprise (New York)1w agoHybrid (New York) · GTM
- GRC Manager1w agoHybrid (San Francisco), Hybrid (Remote) +4 more · EPD
- Software Engineer - Inference Platform2w agoHybrid (San Francisco), Hybrid (Toronto) +3 more · EPD