← All roles

Research Engineer, ML Infrastructure

Cognition· Research & Development
Apply on company site
San FranciscoFull-timePosted 6mo ago

About the role

We are an applied AI lab building end-to-end software agents. Our team builds Devin, the first AI software engineer, and focuses on scalable, reliable systems that enable fast research and training at scale.

Responsibilities

  • Distributed Training Infrastructure: Build and own the systems that run large-scale training jobs reliably across GPU clusters. This includes job launchers, checkpointing and recovery, fault tolerance, and the monitoring that keeps researchers informed and unblocked.
  • Scaling Agent Rollouts: Own the infrastructure that runs hundreds of thousands of concurrent coding agent rollouts in VM sandboxes, from high-fidelity environment design to the distributed systems that hold up at our largest RL training scales.
  • Performance Optimization: Profile and improve training throughput end to end. Identify bottlenecks across data loading, communication overhead, memory utilization, and compute efficiency. Implement solutions that meaningfully improve step time and MFU at scale.
  • Experiment Orchestration and Tooling: Design and maintain the systems researchers use to launch, track, and analyze experiments. Reduce friction in the research loop so that more time is spent on ideas and less on waiting.
  • Data Pipeline Engineering: Build high-throughput, reliable data pipelines for training and evaluation. Ensure data quality, reproducibility, and efficiency at the scale our training runs demand.
  • Debugging and Reliability: Diagnose and resolve training failures across GPUs, networking, numerics, and data. Maintain detailed understanding of failure modes and build systems that fail gracefully and recover fast.
  • Parallelism and Systems Research: Implement and optimize parallelism strategies: data, tensor, pipeline, and sequence parallelism. Understand the tradeoffs deeply and apply them to get the most out of available hardware.
  • Scaling Infrastructure Ahead of Research: Anticipate what the research team will need next and build it before it becomes a constraint. The best infrastructure engineers here are proactive, not reactive.

Qualifications

  • Deep experience building and operating distributed training systems for large models; comfortable owning infrastructure end to end from the cluster level down to the training loop
  • Strong systems engineering fundamentals: distributed systems, networking, storage, and the ability to reason about performance across the full hardware-software stack
  • Proficiency in Python and C++; experience with PyTorch or equivalent deep learning frameworks at a systems level, not just API usage
  • Hands-on experience with GPU performance profiling, memory optimization, and compute efficiency; able to diagnose why a training run is underperforming and fix it
  • Experience implementing or optimizing parallelism strategies (data, tensor, pipeline, sequence) for large model training
  • Track record of building tooling and abstractions that meaningfully accelerate research workflows
  • Strong debugging instincts across complex, distributed systems where failures are non-deterministic and hard to reproduce
  • Enough ML knowledge to engage substantively with researchers: understand what they are training, why the architecture choices matter, and what the infrastructure needs to support
  • We care more about demonstrated capability than credentials. A PhD is one signal among many.

Skills

Required:Distributed TrainingPythonC++PyTorch or equivalentGPU Performance ProfilingParallelism StrategiesSystems DebuggingResearch Collaboration
Last updated Oct 21, 2026. Interested? Apply on the company site.
Cognition logo

About Cognition

cognition is an applied AI lab

An applied AI lab building the future of software engineering. We’re the makers of Devin, a collaborative AI teammate that helps ambitious engineering teams achieve more.

AIai-plg-top-20collaborative AIliquid-ai-recruiting-targetmachine learningsoftware engineering
Company profileWebsite

More open roles at Cognition