C

Causal Labs

Trains physics AI models to forecast weather

Member of Technical Staff - Training Infrastructure

Applied Machine LearningConfirmed live in the last 24 hours
Full-timeSan Francisco, CA, USANo salary listedIn-person

About the Job

Requirements

  • Demonstrated proficiency with distributed training frameworks and techniques, such as Fully Sharded Data Parallel (FSDP), DeepSpeed, Megatron, PyTorch, or JAX/XLA, for training large foundation models.
  • Strong grasp of state-of-the-art techniques for optimizing training workloads, including parallelism strategies, memory optimization, mixed precision, and 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, such as PyTorch and JAX, and their underlying system architectures.

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 do not map cleanly onto existing large language model 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 remain robust under rapid research iteration.
  • Collaborate with researchers to bring novel model architectures from prototype to full scale.
  • Stay up to date on research and bring new ideas to the work.

Desired Qualifications

  • Contributions to open-source machine learning infrastructure, such as PyTorch, Megatron-LM, DeepSpeed, or XLA.

About the Company

Causal Labs

Trains physics AI models to forecast weather

Causal Labs is an AI research company building a foundation model of physical systems for industries whose operations turn on the weather. The company trains on atmospheric sensor and observation data rather than text or images, a domain where cause and effect are governed by physics and yesterday's prediction can be scored against what happened. It runs early pilots with commercial partners and pairs researchers with engineers who deploy the models on customer data. Unlike conventional forecasters, which report what conditions are likely, it is also trying to identify which interventions would change an outcome. The company's goal is to make AI reason about cause and effect by grounding it in the most heavily observed physical system available.

Company size
Not listed
Company stage
Not listed
Total funding
Not listed
Headquarters
Not listed
Founded
Not listed

Recently Posted Jobs

Explore more opportunities at Causal Labs

See more jobs