Skip to content

Repository nanoGPT from @karpathy accelerated with JAX/Flax! The simplest, fastest repository for training/finetuning medium-sized GPTs.

License

Notifications You must be signed in to change notification settings

maxencefaldor/nanoGPT-JAX

Repository files navigation

nanoGPT in JAX !

This repository is based on nanoGPT from karpathy and you will find:

  • nanoGPT.ipynb, the original notebook in PyTorch
  • nanoGPT_jax.ipynb, the original notebook translated in JAX
  • nanoGPT_jax.py, a script to train a nanoGPT of ~200 lines in JAX

You can find an example of the text it generates after training it for ~30 minutes in outputs/output.txt.

Installation

To run this code, you need to clone the repository and install the required libraries with:

git clone ...
pip install -r requirements.txt

Apptainer

To build a container using Apptainer make sure you are in the root of the repository and then run:

apptainer build --fakeroot --force apptainer/container.sif apptainer/container.def

About

Repository nanoGPT from @karpathy accelerated with JAX/Flax! The simplest, fastest repository for training/finetuning medium-sized GPTs.

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published