Research Engineer, Infrastructure, Numerics
About The Position
Thinking Machines is hiring an infrastructure research engineer to design and build core systems for efficient large-scale model training, with a focus on numerics. The role improves the numerical foundations of the distributed training stack, spanning precision formats, kernel optimizations, and communication frameworks to make trillion-parameter model training stable, scalable, and fast. It combines research and systems engineering, including understanding optimization mathematics and the practicalities of distributed compute. This is an evergreen role kept open on an ongoing basis to collect interest. Applications are reviewed continuously, though there may not be an immediate opening that aligns with an applicant’s experience. Applicants may reapply after gaining experience, but should not apply more than once every six months. Separate postings may also be available for specific roles, projects, or teams. The role is based in San Francisco, California. The expected annual salary range is $400,000–$500,000 USD, depending on background, skills, and experience. Thinking Machines sponsors visas and offers health, dental, and vision benefits, unlimited PTO, paid parental leave, and relocation support as needed.
Requirements
- Bachelor’s degree or equivalent experience in computer science, electrical engineering, statistics, machine learning, physics, robotics, or a similar field.
- Understanding of deep learning frameworks such as PyTorch or JAX and their underlying system architectures.
- Ability to thrive in a highly collaborative environment involving cross-functional partners and subject matter experts.
- Initiative and a bias for action, including willingness to work across stacks and teams to help deliver solutions.
- Strong engineering skills and ability to contribute performant, maintainable code and debug complex codebases in areas such as floating-point numerics, low-precision arithmetic, and distributed systems.
Nice To Haves
- Familiarity with distributed frameworks such as PyTorch/XLA, DeepSpeed, or Megatron-LM.
- Experience implementing FP8, INT8, or block-floating point (MX) formats and understanding their numerical trade-offs.
- Prior contributions to open-source deep learning infrastructure such as PyTorch, DeepSpeed, or XLA.
- Publications, patents, or projects related to numerical optimization, communication-efficient training, or systems for large models.
- Experience training and supporting large-scale AI models.
- Track record of improving research productivity through infrastructure design or process improvements.
Responsibilities
- Design and optimize distributed training infrastructure for large-scale LLMs, focusing on performance, stability, and reproducibility across multi-GPU and multi-node setups.
- Implement and evaluate low-precision numerics, such as BF16, MXFP8, and NVFP4, to improve efficiency without sacrificing model quality.
- Develop kernels and communication primitives that use hardware-level support for mixed- and low-precision arithmetic.
- Collaborate with research teams to co-design model architectures and training recipes aligned with emerging numeric formats and stability constraints.
- Prototype and benchmark scaling strategies, including data, tensor, and pipeline parallelism, that integrate precision-adaptive computation and quantized communication.
- Contribute to the design of internal orchestration and monitoring systems so thousands of distributed experiments can run efficiently and reproducibly.
- Share learnings through internal documentation, open-source libraries, or technical reports that advance scalable AI infrastructure.
Benefits
- Health benefits
- Dental benefits
- Vision benefits
- Unlimited PTO
- Paid parental leave
- Relocation support as needed
Stand Out From the Crowd
Upload your resume and get instant feedback on how well it matches this job.
What This Job Offers
Job Type
Full-time
Career Level
Mid Level
