Machine Learning Software Engineer

About The Role

As an ML Engineer on our team, you will work with leaders from industry and academia at the intersection of hardware and software to develop state-of-the-art solutions for emerging problems in AI compute.

The Cerebras software platform is designed to be targeted by today’s most relevant machine learning frameworks such as TensorFlow, PyTorch, and JAX. Our ML software engineers are responsible for integrating these frameworks to work with our own highly optimized software stack. Fundamentally, you will be enabling ML researchers to use the software tools and workflows of today to unlock the advanced hardware capabilities of tomorrow.  

In this role, you will create tools and design workflows that enable the development, training, and deployment of machine learning models on our new hardware system. The workflow covers from a small to an extremely large models with trillion of parameters. Furthermore, you will be lowering abstract computations expressed via third-party ML frameworks into representations that can be compiled into highly optimized executables that target the Cerebras system, and help us design a general backend that can accommodate most advance deep learning models. You will also have the opportunity to participate and contribute to open-source projects that we depend on. 

Responsibilities

• Work on the end-to-end training, eval, and inference workflow with customer-facing API
• Distributed training both via data and model parallelism
• Scale and optimize the data pipeline
• Lower Deep Learning framework graph presentation into our internal IR and add any missing OPs
• Compiler optimization such as graph rewrite, constant folding, common expression elimination, and canonization
• Lower high-level OPs to low-level OPs such as affine dialect
• Handle both static and dynamic computational graph

Requirements

• Bachelor’s, Master’s, or foreign equivalent in Computer Science, Engineering, or related
• 5+ years software development experience
• Understanding of state-of-the-art deep learning model architectures and training protocols
• Strong Python and C++ development skills
• Direct experience with at least one Deep Learning framework internals is strongly preferred

Preferred Skills

• Contributed to a deep learning framework
• Experience with distributed training algorithm
• Familiar with deep learning model architecture
• Experience with MLIR, LLVM, or TVM