{"id":15601002,"url":"https://github.com/lucidrains/triton-transformer","last_synced_at":"2025-04-06T09:06:20.516Z","repository":{"id":46711546,"uuid":"404399082","full_name":"lucidrains/triton-transformer","owner":"lucidrains","description":"Implementation of a Transformer, but completely in Triton","archived":false,"fork":false,"pushed_at":"2022-04-05T00:48:07.000Z","size":35961,"stargazers_count":261,"open_issues_count":5,"forks_count":16,"subscribers_count":15,"default_branch":"main","last_synced_at":"2025-03-30T08:07:38.684Z","etag":null,"topics":["artificial-intelligence","attention-mechanism","deep-learning","gpu-programming","transformers"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lucidrains.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2021-09-08T15:25:58.000Z","updated_at":"2025-03-23T06:43:52.000Z","dependencies_parsed_at":"2022-09-10T06:51:35.484Z","dependency_job_id":null,"html_url":"https://github.com/lucidrains/triton-transformer","commit_stats":null,"previous_names":[],"tags_count":28,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Ftriton-transformer","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Ftriton-transformer/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Ftriton-transformer/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Ftriton-transformer/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/triton-transformer/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247457799,"owners_count":20941906,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["artificial-intelligence","attention-mechanism","deep-learning","gpu-programming","transformers"],"created_at":"2024-10-03T02:11:28.080Z","updated_at":"2025-04-06T09:06:20.497Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"## Transformer in Triton (wip)\n\nImplementation of a Transformer, but completely in \u003ca href=\"https://triton-lang.org/\"\u003eTriton\u003c/a\u003e. I'm completely new to lower-level neural net code, so this repository will mostly be a learning experience, with the end-goal being a vanilla transformer that is faster and more efficient to train.\n\n## Results\n\nLayernorm forward\n\n\u003cimg src=\"./images/layernorm.png\" width=\"400px\"\u003e\u003c/img\u003e\n\nLayernorm forwards and backwards\n\n\u003cimg src=\"./images/layernorm-forward-backward.png\" width=\"400px\"\u003e\u003c/img\u003e\n\nSoftmax forwards and backwards\n\n\u003cimg src=\"./images/softmax.png\" width=\"400px\"\u003e\u003c/img\u003e\n\n## Install\n\n```bash\n$ pip install triton-transformer\n```\n\n## Usage\n\n```python\nimport torch\nfrom triton_transformer import Transformer\n\nmodel = Transformer(\n    num_tokens = 256,       # vocab size\n    max_seq_len = 1024,     # maximum sequence length\n    dim = 512,              # dimension\n    depth = 6,              # depth\n    heads = 8,              # number of heads\n    dim_head = 64,          # dimension per head\n    causal = True,          # autoregressive or not\n    attn_dropout = 0.1,     # attention dropout\n    ff_dropout = 0.1,       # feedforward dropout\n    use_triton = True       # use this to turn on / off triton\n).cuda()\n\nx = torch.randint(0, 256, (1, 1024)).cuda()\nlogits = model(x) # (1, 1024, 256)\n```\n\nTo train, just pass in the labels with the keyword `labels` on forward, and the cross entropy loss will be returned for backprop.\n\nex. BERT\n\n```python\nimport torch\nfrom triton_transformer import Transformer\n\nmodel = Transformer(\n    num_tokens = 20000,\n    max_seq_len = 512,\n    dim = 512,\n    depth = 12,\n    heads = 8,\n    dim_head = 64,\n    use_triton = True\n).cuda()\n\nx = torch.randint(0, 20000, (1, 512)).cuda()\nlabels = torch.randint(0, 20000, (1, 512)).cuda()\nmask = torch.ones(1, 512).bool().cuda()\n\nloss = model(x, mask = mask, labels = labels)\nloss.backward()\n```\n\n## Test - GPT training\n\n```bash\n$ python train.py\n```\n\n## Todo\n\n- [x] softmax\n- [x] cross-entropy (using triton ops)\n- [x] layernorm forward\n- [x] layernorm backwards\n- [x] batch matrix multiply + fused act forwards\n- [x] optimize layernorm backwards (figure out how much to store vs recompute)\n- [x] use memory efficient dropout from Triton tutorials\n- [ ] batch matrix multiply + fused act backwards\n- [ ] fused attention (expand on softmax)\n- [ ] use triton matmul for other projections\n- [ ] benchmark and optimize\n- [ ] kernels conditional on inference vs training\n- [ ] efficient triangular matmul kernel for causal attention\n\n## Citations\n\n```bibtex\n@article{Tillet2019TritonAI,\n    title   = {Triton: an intermediate language and compiler for tiled neural network computations},\n    author  = {Philippe Tillet and H. Kung and D. Cox},\n    journal = {Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages},\n    year    = {2019}\n}\n```\n\n```bibtex\n@misc{vaswani2017attention,\n    title   = {Attention Is All You Need}, \n    author  = {Ashish Vaswani and Noam Shazeer and Niki Parmar and Jakob Uszkoreit and Llion Jones and Aidan N. Gomez and Lukasz Kaiser and Illia Polosukhin},\n    year    = {2017},\n    eprint  = {1706.03762},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CL}\n}\n```\n\n```bibtex\n@misc{so2021primer,\n    title   = {Primer: Searching for Efficient Transformers for Language Modeling},\n    author  = {David R. So and Wojciech Mańke and Hanxiao Liu and Zihang Dai and Noam Shazeer and Quoc V. Le},\n    year    = {2021},\n    eprint  = {2109.08668},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.LG}\n}\n```\n\n```bibtex\n@article{chowdhery2022PaLM,\n  title   = {PaLM: Scaling Language Modeling with Pathways},\n  author  = {Chowdhery, Aakanksha et al},\n  year    = {2022}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Ftriton-transformer","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Ftriton-transformer","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Ftriton-transformer/lists"}