Search
Go
Research Engineer, Infrastructure, Numerics
Summary by AIEngineer.careers
Research Engineer, Infrastructure, Numerics is a AI Infrastructure Engineer role (full-time). with Thinking Machines Lab. in SAN FRANCISCO, US. Compensation shown: $350K–$475K. Imported listing (source: zerogtalent.com). Apply on the employer's site (zerogtalent.com).
Imported listing (source: zerogtalent.com) · Apply on zerogtalent.com
SAN FRANCISCO · US
$350K–$475K
Full-time
SENIOR
The sections below reproduce the third-party job description for reference. AIEngineer.careers does not write or control this text.
Imported job description
Sourced from zerogtalent.com
Thinking Machines Lab is seeking an AI Infrastructure Engineer to design and optimize distributed training systems for large-scale language models. You will focus on numerical foundations, including low-precision formats, kernel optimization, and communication primitives, to make trillion-parameter model training stable and efficient. This role bridges research and systems engineering, requiring deep expertise in deep learning frameworks and distributed compute architectures.
Responsibilities
- Design and optimize distributed training infrastructure for multi-GPU and multi-node LLM setups
- Implement and evaluate low-precision numerics (BF16, MXFP8, NVFP4) for efficiency
- Develop hardware-level kernels and communication primitives for mixed-precision arithmetic
- Collaborate with research teams to co-design model architectures and training recipes
- Prototype and benchmark scaling strategies including data, tensor, and pipeline parallelism
- Contribute to internal orchestration and monitoring systems for scalable experiments
Requirements
- Bachelor’s degree or equivalent in CS, EE, Statistics, ML, Physics, or Robotics
- Strong understanding of deep learning frameworks (PyTorch, JAX) and system architectures
Thinking Machines Lab
San Francisco · US · 135+ employees
Thinking Machines Lab is an American artificial intelligence startup founded by former OpenAI CTO Mira Murati. The company builds open-weights large language models and provides research infrastructure to make AI more customizable and accessible for unique needs in scientific and programming domains.
