Research Engineer, Infrastructure, Numerics

Thinking Machines Lab
San Francisco
Workplace: HybridFull timeUSD 350,000 - 475,000 annuallyFunction: Research & Scientific (R&D)Education: bachelorsSkills: ["Collaboration","Initiative","Bias for action","Debugging","Problem-solving"]

Design and optimize distributed training infrastructure for large-scale LLMs with a focus on numerics. You’ll implement and evaluate low-precision formats, build kernels and communication primitives for mixed/low-precision arithmetic, and prototype scaling strategies across multi-GPU, multi-node setups. Collaborate with research teams to co-design training recipes, and contribute to internal orchestration and monitoring so thousands of experiments run efficiently, stably, and reproducibly.

Loading

Loading job details...

Preparing the role view and application actions.

FursaFursa
Thinking Machines Lab
Thinking Machines Lab
1 day ago

Research Engineer, Infrastructure, Numerics

✓ Verified Job

Canonical indexed version, validated from employer's careers page.

Source: Company careers pageValidated by: Fursa AI
Last checked: 9 hours agoStatus: Live

Job Summary

Design and optimize distributed training infrastructure for large-scale LLMs with a focus on numerics. You’ll implement and evaluate low-precision formats, build kernels and communication primitives for mixed/low-precision arithmetic, and prototype scaling strategies across multi-GPU, multi-node setups. Collaborate with research teams to co-design training recipes, and contribute to internal orchestration and monitoring so thousands of experiments run efficiently, stably, and reproducibly.
Location: San Francisco
Workplace: Hybrid
Employment Type: Full time
Job Function: Research & Scientific (R&D)

Key Responsibilities

  • •Design and optimize distributed training infrastructure for large-scale LLMs, emphasizing performance, stability, and reproducibility.
  • •Implement and evaluate low-precision numerics (e.g., BF16, MXFP8, NVFP4) to improve efficiency without degrading model quality.
  • •Develop kernels and communication primitives that leverage hardware 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.
  • •Contribute to internal orchestration and monitoring systems so thousands of distributed experiments run efficiently and reproducibly.

Pay and Benefits

Salary: USD 350,000 - 475,000 annually
Perks:Health InsuranceDentalVisionPaid LeaveRelocation

Key Requirements

  • •Bachelor’s degree (or equivalent experience) in computer science, electrical engineering, statistics, machine learning, physics, robotics, or a related field.
  • •Understanding of deep learning frameworks (e.g., PyTorch, JAX) and their underlying system architectures.
  • •Strong engineering skills to contribute performant, maintainable code and debug complex codebases in areas like floating-point numerics, low-precision arithmetic, and distributed systems.
  • •Ability to thrive in a highly collaborative, cross-functional environment with diverse subject matter experts.
  • •A bias for action and initiative to work across stacks and teams to ensure key work ships.
Experience:Machine learningDeep learningDistributed training
Education:Bachelor's
Skills:CollaborationInitiativeBias for actionDebuggingProblem-solving
Tech Stack:PyTorchJAXPyTorch/XLADeepSpeedMegatron-LMBF16FP8INT8MXFP8NVFP4MXFloating-point numericsLow-precision arithmeticDistributed systemsMulti-GPUMulti-nodeTensor parallelismPipeline parallelismQuantized communication

Eligibility

Work Authorization:Sponsorship available.

Company Brief

Thinking Machines Lab
Develops enterprise AI solutions, custom large language models, and ML platforms to help organizations deploy intelligent applications. Services include data engineering, model development, and AI consulting for scale and production readiness.
Industry: AI & Machine Learning
Company Size: Medium (51 to 250 employees)
Growth: Growth Stage Startup
Headquarters: Mumbai, India
Website