← All roles

Research Engineer, Infrastructure, Numerics

Thinking Machines Lab· Research Infrastructure (ML Infra)
Apply on company site
San Francisco, CA$350k - $475kFull-timePosted 11mo ago

About the role

Infrastructure research engineer focusing on numerics to design and build core systems for efficient large-scale model training, with emphasis on precision formats, kernel optimizations, and communication frameworks for stable, scalable trillion-parameter models.

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 (e.g., BF16, MXFP8, 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 numeric formats and stability constraints.
  • Prototype and benchmark scaling strategies such as data, tensor, and pipeline parallelism that integrate precision-adaptive computation and quantized communication.
  • Contribute to the design of internal orchestration and monitoring systems to run thousands of distributed experiments efficiently and reproducibly.
  • Publish and share learnings through internal documentation, open-source libraries, or technical reports that advance scalable AI infrastructure.

Qualifications

  • Bachelor’s degree or equivalent experience in computer science, electrical engineering, statistics, machine learning, physics, robotics, or similar.
  • Understanding of deep learning frameworks (e.g., PyTorch, JAX) and their underlying system architectures.
  • Thrive in a highly collaborative environment involving many cross-functional partners and subject matter experts.
  • Bias for action with initiative to work across stacks and teams to ship—strong engineering skills to contribute performant, maintainable code and debug in complex codebases (floating-point numerics, low-precision arithmetic, distributed systems).

Skills

Required:Distributed training infrastructureLow-precision numericsKernels and communication primitives

Benefits

  • Generous health, dental, and vision benefits
  • Unlimited PTO
  • Paid parental leave
  • Relocation support as needed
Last updated Jul 8, 2026. Interested? Apply on the company site.
Thinking Machines Lab logo

About Thinking Machines Lab

Connectionism: Research Blog by Thinking Machines Lab

Thinking Machines Lab is an artificial intelligence research and product company. We’re building a future where everyone has access to the knowledge and tools to make AI work for their unique needs and goals.

AIartificial intelligencebessemer-ai-infrastructure-2026deep learningliquid-ai-recruiting-targetmachine learningproductreinforcement-learning-platforms
Company profileWebsite

More open roles at Thinking Machines Lab