Research Engineer (LLM Performance), London

Requirements

• Significant experience with large scale distributed training of LLMs.
• Experience with deep learning ML frameworks (either JAX or PyTorch).
• Knowledge of parallelism strategies and collective communication libraries (e.g. NCCL).
• Good understanding of GPU architectures. Reasoning about performance concepts is more important than writing kernels from scratch
• Excellent collaboration skills.

Nice to haves

• Experience with general LLM serving stacks.
• Knowledge of XLA, Triton, Pallas, CUDA or similar accelerator DSLs / compilers.
• Experience with optimising ML accuracy using low-precision formats.
• Prior experience building, deploying and maintaining production systems on GCP.
• Interest in chemistry and biology.

What you'll be doing

• Implement and optimize LLM post-training methods at scale on frontier models.
• Collaborate with research teams to translate new methods into production-ready systems.
• Relentlessly prioritize and execute on performance optimization opportunities.
• Evaluate and deploy frameworks for supervised fine-tuning, reinforcement learning and LLM evaluation.
• Diagnose and fix performance bottlenecks and communication overhead in distributed training and inference systems.
• Deploy low-precision methods to balance performance with accuracy, impacting real-world drug design programs.

Perks and Benefits

• Hybrid working model (3 days in the office per week).
• Commitment to equal employment opportunities.
• Culture emphasizing collaboration, learning, and support.