# Research Engineer, Infrastructure, Numerics

**Company:** [Thinking Machines Lab](https://hotfix.jobs/companies/thinking-machines-lab)
**Location:** San Francisco, CA
**Role:** DevOps / SRE
**Salary:** $350k – $475k/yr
**Skills:** PyTorch, JAX, Deepspeed, Megatron-Lm, Pytorch/Xla, Xla, Bf16, Fp8, Int8, Distributed Systems
**Posted:** 2026-05-04

> Designs and optimizes distributed training infrastructure for large-scale LLMs, focusing on low-precision numerics, kernel optimizations, and communication frameworks to enable stable, scalable trillion-parameter model training. Requires strong systems engineering, deep learning frameworks knowledge, and collaborative research mindset.

## Job Description

## What You’ll Do

- 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 (for example, **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 that align with emerging 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 our internal orchestration and monitoring systems to ensure that thousands of distributed experiments can run efficiently and reproducibly.
- Publish and share learnings through internal documentation, open-source libraries, or technical reports that advance the field of scalable AI infrastructure.

## Skills and Qualifications

**Minimum 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, different cross-functional partners and subject matter experts.
- A bias for action with a mindset to take initiative to work across different stacks and different teams where you spot the opportunity to make sure something ships.
- Strong engineering skills, ability to contribute performant, maintainable code and debug in complex codebases in areas such as floating-point numerics, low-precision arithmetic, and distributed systems.

**Preferred qualifications:**
- Familiarity with distributed frameworks such as **PyTorch/XLA**, **DeepSpeed**, **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.

## Logistics

**Compensation:** Depending on background, skills and experience, the expected annual salary range for this position is **$350,000 - $475,000 USD**.

**Benefits:** Generous health, dental, and vision benefits, unlimited PTO, paid parental leave, and relocation support as needed.

## Similar roles

- [Reliability Engineer, Supercomputing](https://hotfix.jobs/jobs/e9263e4b-66f2-4203-b8da-b88ba1855dfe) - Thinking Machines Lab - San Francisco, CA - $350k – $475k/yr
- [Network Engineer, Supercomputing](https://hotfix.jobs/jobs/2a952a8c-f34c-4fdd-a219-c8ac22399145) - Thinking Machines Lab - San Francisco, CA - $350k – $475k/yr
- [Performance Engineer, Inference Systems](https://hotfix.jobs/jobs/c4e8b3e8-25ad-4a83-b11d-bebf411d63f1) - Anthropic - San Francisco, CA - $350k – $850k/yr
- [Site Reliability Engineer (SRE)](https://hotfix.jobs/jobs/893a897e-0421-40b8-a8ac-f61c0b30d5ff) - Thinking Machines Lab - San Francisco, CA - $350k – $475k/yr
- [Research Engineer, Infrastructure, Training Systems](https://hotfix.jobs/jobs/9cc8be27-572e-4c19-97d0-8a7ae54e41a8) - Thinking Machines Lab - San Francisco, CA - $350k – $475k/yr

**Apply:** https://hotfix.jobs/jobs/446d787b-0944-43fb-81c5-5783517b66fb
**Canonical:** https://hotfix.jobs/jobs/446d787b-0944-43fb-81c5-5783517b66fb