Jobs Companies Causal Labs Member of Technical Staff — Training Infrastructure

About this Member of Technical Staff — Training Infrastructure role at Causal Labs

Causal Labs · Onsite · San Francisco

Our mission is general causal intelligence; AI that is capable of (1) predicting the future and (2) identifying the actions to alter it.

To achieve this breakthrough, we are building a Large Physics foundation Model (LPM) because physical systems, unlike text or images, are governed by verifiable cause and effect. We believe that scaling on physics will enable an understanding of causality required to predict and control physical systems, starting with weather.

Our founding team has built and deployed AI against the physical world in robotics, drug discovery, and particle physics at institutions like DeepMind, Waymo, Cruise, Insitro, Nabla Bio, and CERN.

 

We look for infrastructure engineers who are excited to tackle unsolved problems. Training an LPM means scaling novel architectures over multimodal physical data — a problem where the playbooks from language and vision only partially apply. Your mission is to make large-scale training fast, efficient, and reliable, so that every GPU cycle accelerates research progress.

Responsibilities

  • Design, implement, and optimize distributed training systems that scale across thousands of GPUs

  • Research and test parallelization strategies and numerical precision trade-offs across model scales, including for architectures that don't map cleanly onto existing LLM training stacks

  • Analyze, profile, and debug low-level GPU operations to maximize throughput and hardware utilization

  • Build reusable frameworks for checkpointing, fault tolerance, and reproducibility that stay robust under rapid research iteration

  • Collaborate with researchers to bring novel model architectures from prototype to full scale

  • Stay up-to-date on research to bring new ideas to work

What we're looking for

We value a relentless approach to problem-solving, rapid execution, and the ability to quickly learn in unfamiliar domains.

  • Demonstrated proficiency with distributed training frameworks and techniques (e.g. FSDP, DeepSpeed, Megatron, Pytorch, JAX/XLA) to train large foundation models

  • Strong grasp of state-of-the-art techniques for optimizing training workloads: parallelism strategies, memory optimization, mixed precision, communication overlap

  • Ability to profile and debug performance in complex codebases, from framework internals down to kernels and collectives

  • Deep understanding of deep learning frameworks (e.g. PyTorch, JAX) and their underlying system architectures

  • Bonus: contributions to open-source ML infrastructure (e.g. PyTorch, Megatron-LM, DeepSpeed, XLA)

Ready to apply to Causal Labs?
Apply to Causal Labs

Similar jobs

Sign up for suggestions tailored to the jobs you open and the searches you save.

More jobs at Causal Labs

See all jobs at Causal Labs →

Apply now
🤖

Whoa — hold up

JobsRadar was built for real people having a rough time in their job search — not for automated requests. You're clicking way too fast and you're now temporarily blocked.

Come back later. If you're genuinely job hunting, we've got your back — just act like a human.

Catch your next role the second it’s posted.

Create a free account and we’ll watch the boards for you — the instant a job matches your search, it lands in your inbox or Telegram. No digging, no refreshing.

Create free account

Free forever · takes 30 seconds · already have one?

Get an edge on your job hunt.

Join our Telegram channel for the stuff that helps you land the role — salary benchmarks, the weekly market pulse, and new-feature drops. No spam, just signal.

Join the channel — it's free