https://github.com/jaykef/triton-nanogpt

https://github.com/jaykef/triton-nanogpt

Last synced: 7 days ago
JSON representation

• Host: GitHub
• URL: https://github.com/jaykef/triton-nanogpt
• Owner: Jaykef
• Created: 2024-09-25T09:24:18.000Z (about 1 month ago)
• Default Branch: main
• Last Pushed: 2024-10-21T01:03:51.000Z (16 days ago)
• Last Synced: 2024-10-21T04:29:36.759Z (15 days ago)
• Language: Python
• Size: 309 KB
• Stars: 12
• Watchers: 1
• Forks: 0
• Open Issues: 0
• Metadata Files:
  • Readme: README.md

Awesome Lists containing this project

README

        
# Triton-Accelerated NanoGPT

The WHY behind this ordeal

After practicing triton for about 2 weeks now, I challenged myself into implementing custom triton kernels for Karpathy's nanoGPT and quite an ordeal it was but somehow got something working, not perfect but getting there:), contributions are welcomed.

## Kernels

Supports lightweight custom triton kernels for softmax, layer normalization, cross entropy loss and GELU activation.

## Training

GPU-aware train loop with effective gradient accumulation, learning rate scheduling and gradient clipping with val loss tracking.

- **Setup**: Requires GPU! Ensure you have PyTorch and Triton installed. GPU Poor? I am too, I used one free T4 on google colab.

- **Data**: Using Tiny Shakespeare dataset by default. It will be downloaded automatically if not present.

- **Training**: 

   ```python

   python triton_nanoGPT.py

   ```

This will train for 100 epochs, save checkpoint as `nanoGPT_cpkt.pth` and sample from it.

## License

MIT