History

Angelos Katharopoulos e9b32747b4 Add grad checkpointing and PE in the transformer example (#387 ) * Add grad checkpointing and PE in the transformer example * Remove other frameworks from LM example * Remove the other frameworks from MNIST example * Improve the transformer LM example * Fix black and change LR		2024-02-01 13:04:03 -08:00
..
datasets.py	Add llms subdir + update README (#145 )	2023-12-20 10:22:25 -08:00
main.py	Add grad checkpointing and PE in the transformer example (#387 )	2024-02-01 13:04:03 -08:00
README.md	Add grad checkpointing and PE in the transformer example (#387 )	2024-02-01 13:04:03 -08:00
requirements.txt	Add grad checkpointing and PE in the transformer example (#387 )	2024-02-01 13:04:03 -08:00

Transformer LM

This is an example of a decoder-only Transformer LM. The only dependency is MLX.

Run the example on the GPU with:

python main.py --gpu

By default the dataset is the PTB corpus. Choose a different dataset with the --dataset option.