Principal Staff Software Engineer, AI Training Platform

This role will be based in Mountain View, CA.

At LinkedIn, we trust each other to do our best work where it works best for us and our teams. This role offers hybrid work options, meaning you can work from home and commute to a LinkedIn office, depending on what's best for you and when your team needs to be together.

As part of LinkedIn's AI Platform group, the AI Training team is responsible for developing and maintaining highly available and scalable deep learning training solutions to power our rapidly growing AI use cases. The team is responsible for scaling LinkedIn's AI model training with hundreds of billions of parameters for all AI use cases from recommendation models, large language models (Generative AI), to computer vision models. We optimize training performance across algorithms, AI frameworks, infrastructure software, and hardware to harness the power of our GPU fleet with thousands of latest GPU cards. The team also works closely with the open source community and has many open source committers (TensorFlow, Horovod, Ray, Hadoop, etc.) in the team. Additionally, this team focussed on technologies like LLMs, GNNs, Incremental Learning, Online Learning, and advanced LLM Agents work for Training infrastructure.

As a Principal Staff Software Engineer on the AI Training Infra team, you will play a crucial role in leading and building the next-gen training infrastructure to power AI use cases. You will design and implement high performance AI Training pipeline, data I/O, work with open source teams to identify and resolve issues in popular libraries like Huggingface, Horovod and PyTorch, debug and optimize deep learning training, and provide advanced support for internal AI teams in areas like model parallelism, data parallelism, Zero, automatic mixed precision and kernel fusion. Finally, you will assist in and guide the development of containerized pipeline orchestration infrastructure, including developing and distributing stable base container images, providing advanced profiling and observability, and updating internally maintained versions of deep learning frameworks and their companion libraries like Tensorflow, PyTorch, DeepSpeed, GNNs, Flash Attention and more.

Responsibilities

• Owning the technical strategy for broad or complex requirements with insightful and forward-looking approaches that go beyond the direct team and solve large open-ended problems.
• Designing, implementing, and optimizing the performance of large-scale distributed training for personalized recommendation as well as large language models.
• Improving the observability and understandability of various systems with a focus on improving developer productivity and system sustenance.
• Mentoring other engineers, defining our challenging technical culture, and helping to build a fast-growing team.
• Working closely with the open-source community to participate and influence cutting edge open-source projects (e.g., PyTorch, GNNs, DeepSpeed, Huggingface, etc.).
• Functioning as the tech-lead for several concurrent key initiatives for the Training Infrastructure and defining the future of AI training platforms.