Build a Mini GPT with JAX

JAX is an open-source machine learning library. In this worksop, we'll walk you through how to write and train a "small" large language model using JAX. You'll work through a notebook that contains an implementation of a mini GPT, and we'll explain concepts as we go (from data loading, to basic transformers, to distributed training, to inference).

Please bring a laptop. We'll use Google Colab to run our example. There is nothing to install in advance. If you'd like to get the flavor of the code we'll explore, you can check out this tutorial now.

Knowledge prerequisites: Basic knowledge of Python 

Hardware/software prerequisites: Bring a laptop with a web browser.

Workshop format: Demonstration and hands-on

Target audience: Students, researchers, faculty, staff

See the full Research Computing training schedule or subscribe to the Research Computing mailing list.

Research Computing | View More Events
Co-hosted with: GradFUTURES

Contact the organizers