Research Engineer, Infrastructure, Training Systems

Thinking Machines Lab•San Francisco, CA

19h

About The Position

Thinking Machines is hiring an engineer to build and improve the systems that support large-scale frontier-model pre-training and mid-training. The role focuses on keeping large accelerator clusters efficient and reliable during training runs that may last days or weeks. Work spans the training stack, including model partitioning across accelerators, numerical and hardware efficiency, distributed workload coordination, and recovery from failures. The position may focus on distributed training—such as frameworks, runtimes, sharded model and optimizer states, parallelism, checkpointing, scheduling, and model FLOPs utilization—or on communication and collectives, including accelerator communication runtimes, collective algorithms, topology-aware communication, communication/computation overlap, and fault tolerance. Specialization in both areas is not required. The training systems stack is written in Rust; prior Rust experience is helpful, and strong systems engineers interested in learning Rust are encouraged to apply. The engineer will collaborate closely with researchers to understand workloads, develop and test systems changes, and measure their effect on performance and correctness. The role also involves building reusable, maintainable systems that improve experimentation and reproducibility, and sharing findings through documentation, open-source contributions, or technical reports. Focus will depend on the candidate’s strengths and the team’s needs.

Requirements

Strong engineering skills, including writing performant, maintainable code and debugging complex systems.
Hands-on experience building, optimizing, or operating distributed model-training infrastructure, with a clear understanding of personal contributions.
Depth in at least one relevant area, such as distributed training frameworks, parallelism, GPU communication, low-precision training, runtime systems, or training-workload orchestration.
Understanding of deep learning frameworks such as PyTorch or JAX and the systems beneath them.
Ability to investigate problems across subsystem boundaries and independently carry a solution through implementation and validation.
Collaborative approach to working with researchers and engineers, including clear communication about evidence, uncertainty, and tradeoffs.
Bachelor’s degree in a relevant technical field or equivalent practical experience.

Nice To Haves

Experience implementing or improving FSDP, ZeRO, tensor, pipeline, expert, context, or sequence parallelism.
Experience improving training MFU, throughput, memory efficiency, or cost through profiling, mixed precision, FP4/FP8, kernels, or compiler work.
Experience developing NCCL or equivalent communication runtimes, collective algorithms, or topology-aware communication.
Experience improving distributed checkpointing, restart behavior, fault tolerance, or training stability.
Experience building scheduling or runtime infrastructure for large, long-running model-training workloads.
Substantial contributions to projects such as PyTorch, XLA, Megatron-LM, DeepSpeed, or related infrastructure.
Prior Rust experience.

Responsibilities

Design, implement, and operate distributed training infrastructure that scales across thousands of GPUs and nodes.
Improve training throughput, accelerator utilization, and time to train by diagnosing bottlenecks across compute, memory, networking, and storage.
Develop parallelism strategies, runtime components, communication primitives, or workload orchestration for evolving model architectures.
Improve the reliability of long-running training jobs through observability, checkpointing, failure diagnosis, and recovery.
Work with researchers to understand new workloads, test systems changes, and measure effects on training performance and correctness.
Build reusable, maintainable frameworks and libraries that make experimentation faster and training more reproducible.
Share useful findings through documentation, open-source contributions, or technical reports.