Research Engineer (LLM Performance), London

Your impact 

This is an exciting opportunity for you to contribute to frontier research at the intersection of AI and drug design.

Working in a highly creative, iterative environment, you will join the model performance and scaling team, where you will partner with scientists and engineers to scale foundational models that will transform the biopharmaceutical world as we know it.  

You will draw upon your existing engineering experience whilst learning from those around you, to apply novel techniques and ideas to newly encountered model and systems performance optimizations, as well as machine learning, computational biology and chemistry problems.

What you will do 

• 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.

Skills and qualifications 

Essential:

• 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 have:

• 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.